diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index bdfd26a79b..1afe119524 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -308,7 +308,8 @@ impl ConnectionConfig { format!("{scheme}://{host}:{port}/") } DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks => { - format!("mysql://{host}:{port}{db_part}?{params}") + let suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; + format!("mysql://{host}:{port}{db_part}{suffix}") } DatabaseType::Postgres | DatabaseType::Redshift => { let suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; @@ -405,7 +406,8 @@ impl ConnectionConfig { } } DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks => { - format!("mysql://{}:{}@{host}:{port}{db_part}?{params}", username, password) + let suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; + format!("mysql://{}:{}@{host}:{port}{db_part}{suffix}", username, password) } DatabaseType::Postgres | DatabaseType::Redshift => { let suffix = if params.is_empty() { String::new() } else { format!("?{params}") }; @@ -548,31 +550,11 @@ impl ConnectionConfig { fn normalized_url_params(&self) -> String { let value = self.url_params.as_deref().unwrap_or("").trim(); if self.needs_bare_mysql() { - let v = value.trim_start_matches('?'); - let filtered: Vec<&str> = v - .split('&') - .filter(|p| !p.is_empty() && !p.starts_with("charset=") && !p.starts_with("ssl-mode=preferred")) - .collect(); - return if filtered.is_empty() { - "ssl-mode=disabled".to_string() - } else { - format!("ssl-mode=disabled&{}", filtered.join("&")) - }; + return normalize_bare_mysql_url_params(value); } match self.db_type { DatabaseType::Mysql => normalize_mysql_url_params(value, self.ssl, self.ca_cert_path.trim().is_empty()), - DatabaseType::Doris | DatabaseType::StarRocks => { - let v = value.trim_start_matches('?'); - let filtered: Vec<&str> = v - .split('&') - .filter(|p| !p.is_empty() && !p.starts_with("charset=") && !p.starts_with("ssl-mode=preferred")) - .collect(); - if filtered.is_empty() { - "ssl-mode=disabled".to_string() - } else { - format!("ssl-mode=disabled&{}", filtered.join("&")) - } - } + DatabaseType::Doris | DatabaseType::StarRocks => normalize_bare_mysql_url_params(value), DatabaseType::Postgres | DatabaseType::Redshift => normalize_postgres_url_params(value, self.ssl), DatabaseType::MongoDb => value.trim_start_matches('?').to_string(), _ => value.trim_start_matches('?').to_string(), @@ -590,6 +572,23 @@ fn url_params_contains_flag(params: Option<&str>, key: &str, expected: &str) -> ) } +fn normalize_bare_mysql_url_params(value: &str) -> String { + value + .trim_start_matches('?') + .split('&') + .filter(|part| { + !part.is_empty() + && !url_param_key_is(part, "charset") + && !url_param_key_is(part, "ssl-mode") + && !url_param_key_is(part, "sslmode") + && !url_param_key_is(part, "require_ssl") + && !url_param_key_is(part, "verify_ca") + && !url_param_key_is(part, "verify_identity") + }) + .collect::>() + .join("&") +} + fn normalize_mysql_url_params(value: &str, force_tls: bool, accept_invalid_certs: bool) -> String { let value = value.trim_start_matches('?'); let mut parts: Vec = value.split('&').filter(|part| !part.is_empty()).map(str::to_string).collect(); @@ -1022,7 +1021,32 @@ mod tests { config.driver_profile = Some("oceanbase".to_string()); assert!(config.needs_bare_mysql()); - assert_eq!(config.connection_url(), "mysql://user%40tenant%23cluster:secret@10.1.2.3:2883?ssl-mode=disabled"); + assert_eq!(config.connection_url(), "mysql://user%40tenant%23cluster:secret@10.1.2.3:2883"); + } + + #[test] + fn starrocks_profile_omits_mysql_ssl_mode_param() { + let mut config = mysql_config("root", "secret", Some("analytics")); + config.driver_profile = Some("starrocks".to_string()); + config.url_params = Some( + "ssl-mode=disabled&sslmode=required&require_ssl=true&verify_ca=false&verify_identity=false&charset=utf8mb4" + .to_string(), + ); + + assert!(config.needs_bare_mysql()); + assert_eq!(config.connection_url(), "mysql://root:secret@10.1.2.3:2883/analytics"); + } + + #[test] + fn starrocks_profile_keeps_non_mysql_tls_params() { + let mut config = mysql_config("root", "secret", Some("analytics")); + config.driver_profile = Some("starrocks".to_string()); + config.url_params = Some("connect_timeout=10&sessionVariables=query_timeout=60".to_string()); + + assert_eq!( + config.connection_url(), + "mysql://root:secret@10.1.2.3:2883/analytics?connect_timeout=10&sessionVariables=query_timeout=60" + ); } #[test] diff --git a/crates/dbx-core/src/object_source_sql.rs b/crates/dbx-core/src/object_source_sql.rs index ef79d7b6b6..a1868bbf9a 100644 --- a/crates/dbx-core/src/object_source_sql.rs +++ b/crates/dbx-core/src/object_source_sql.rs @@ -125,7 +125,7 @@ pub fn build_executable_object_source_statements(input: EditableObjectSourceSqlI return Ok(vec![replace_sqlserver_create_with_alter(source)]); } - if matches!(input.database_type, DatabaseType::Postgres | DatabaseType::Gaussdb) + if matches!(input.database_type, DatabaseType::Postgres | DatabaseType::Gaussdb | DatabaseType::OpenGauss) && input.object_type == ObjectSourceKind::View { return Ok(vec![format!( @@ -211,6 +211,7 @@ fn is_postgres_like(database_type: DatabaseType) -> bool { DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb + | DatabaseType::OpenGauss | DatabaseType::Kingbase | DatabaseType::Highgo | DatabaseType::Vastbase @@ -423,6 +424,22 @@ mod tests { ); } + #[test] + fn opengauss_view_body_opens_as_create_or_replace_view() { + let sql = build_executable_object_source_sql(EditableObjectSourceSqlInput { + database_type: DatabaseType::OpenGauss, + object_type: ObjectSourceKind::View, + schema: Some("public".to_string()), + name: "active users".to_string(), + source: " SELECT id, name FROM users WHERE active ".to_string(), + }) + .unwrap(); + assert_eq!( + sql, + "CREATE OR REPLACE VIEW \"public\".\"active users\" AS\nSELECT id, name FROM users WHERE active;" + ); + } + #[test] fn view_ddl_wraps_postgres_body_as_create_or_replace_view() { let sql = build_view_ddl_sql(BuildViewDdlInput { diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 1a5bc0ae71..7af622ec25 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1,6 +1,6 @@ use crate::connection::{AppState, MysqlMode, PoolKind}; use crate::db; -use crate::models::connection::DatabaseType; +use crate::models::connection::{ConnectionConfig, DatabaseType}; use std::future::Future; macro_rules! extract_pool { @@ -805,15 +805,31 @@ pub async fn get_table_ddl_core( let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Pool not found")?; + let db_config = connection_config(state, connection_id).await; match pool { PoolKind::Mysql(p, _) => mysql_ddl(p, table).await, + PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_opengauss_family_config) => { + match opengauss_table_ddl(p, schema, table).await { + Ok(ddl) => Ok(ddl), + Err(_) => pg_ddl(p, schema, table).await, + } + } PoolKind::Postgres(p) => pg_ddl(p, schema, table).await, PoolKind::Sqlite(p) => sqlite_ddl(p, table).await, _ => Err("DDL not supported for this database type".to_string()), } } +async fn connection_config(state: &AppState, connection_id: &str) -> Option { + state.configs.read().await.get(connection_id).cloned() +} + +fn is_opengauss_family_config(config: &ConnectionConfig) -> bool { + matches!(config.db_type, DatabaseType::OpenGauss | DatabaseType::Gaussdb) + || matches!(config.driver_profile.as_deref(), Some("opengauss" | "gaussdb")) +} + fn sql_string(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } @@ -856,7 +872,15 @@ pub fn sqlserver_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSo pub fn postgres_object_source_sql(schema: &str, name: &str, kind: &db::ObjectSourceKind) -> String { match kind { db::ObjectSourceKind::View => { - format!("SELECT pg_get_viewdef('{}.{}'::regclass, true)", pg_ident(schema), pg_ident(name)) + format!( + "SELECT pg_get_viewdef(c.oid, 0) \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = {} AND c.relname = {} AND c.relkind IN ('v','m') \ + ORDER BY c.oid LIMIT 1", + sql_string(schema), + sql_string(name) + ) } db::ObjectSourceKind::Procedure | db::ObjectSourceKind::Function => { let prokind = if matches!(kind, db::ObjectSourceKind::Procedure) { "p" } else { "f" }; @@ -904,6 +928,17 @@ pub fn mysql_object_source_sql(name: &str, kind: &db::ObjectSourceKind) -> Strin } } +pub fn postgres_view_source_fallback_sql(schema: &str, name: &str) -> String { + format!( + "SELECT definition \ + FROM pg_catalog.pg_views \ + WHERE schemaname = {} AND viewname = {} \ + LIMIT 1", + sql_string(schema), + sql_string(name) + ) +} + fn first_string_cell(result: db::QueryResult) -> Result { result .rows @@ -961,9 +996,7 @@ pub async fn get_object_source_core( } else { match connections.get(&pool_key).ok_or("Pool not found")? { PoolKind::Mysql(pool, _) => mysql_object_source(pool, name, &object_type).await?, - PoolKind::Postgres(pool) => first_string_cell( - db::postgres::execute_query(pool, &postgres_object_source_sql(schema, name, &object_type)).await?, - )?, + PoolKind::Postgres(pool) => postgres_object_source(pool, schema, name, &object_type).await?, PoolKind::Sqlite(pool) => first_string_cell( db::sqlite::execute_query(pool, &sqlite_object_source_sql(name, &object_type)).await?, )?, @@ -989,6 +1022,26 @@ pub async fn get_object_source_core( }) } +async fn postgres_object_source( + pool: &deadpool_postgres::Pool, + schema: &str, + name: &str, + object_type: &db::ObjectSourceKind, +) -> Result { + let sql = postgres_object_source_sql(schema, name, object_type); + match db::postgres::execute_query(pool, &sql).await.and_then(first_string_cell) { + Ok(source) => Ok(source), + Err(primary_err) if matches!(object_type, db::ObjectSourceKind::View) => { + let fallback_sql = postgres_view_source_fallback_sql(schema, name); + db::postgres::execute_query(pool, &fallback_sql) + .await + .and_then(first_string_cell) + .map_err(|fallback_err| format!("{primary_err}; fallback failed: {fallback_err}")) + } + Err(err) => Err(err), + } +} + #[cfg(test)] mod object_source_tests { use super::*; @@ -1006,7 +1059,7 @@ mod object_source_tests { fn builds_postgres_object_source_sql_for_views_and_functions() { assert_eq!( postgres_object_source_sql("public", "active_users", &ObjectSourceKind::View), - "SELECT pg_get_viewdef('\"public\".\"active_users\"'::regclass, true)" + "SELECT pg_get_viewdef(c.oid, 0) FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = 'public' AND c.relname = 'active_users' AND c.relkind IN ('v','m') ORDER BY c.oid LIMIT 1" ); assert_eq!( postgres_object_source_sql("public", "recalc_score", &ObjectSourceKind::Function), @@ -1014,6 +1067,25 @@ mod object_source_tests { ); } + #[test] + fn builds_postgres_view_source_sql_without_regclass_cast() { + let sql = postgres_object_source_sql("tenant's schema", "active users", &ObjectSourceKind::View); + + assert!(!sql.contains("::regclass")); + assert!(sql.contains("pg_get_viewdef(c.oid, 0)")); + assert!(sql.contains("n.nspname = 'tenant''s schema'")); + assert!(sql.contains("c.relname = 'active users'")); + assert!(sql.contains("c.relkind IN ('v','m')")); + } + + #[test] + fn builds_postgres_view_source_fallback_sql_from_pg_views() { + assert_eq!( + postgres_view_source_fallback_sql("tenant's schema", "active users"), + "SELECT definition FROM pg_catalog.pg_views WHERE schemaname = 'tenant''s schema' AND viewname = 'active users' LIMIT 1" + ); + } + #[test] fn builds_oracle_object_source_sql_using_metadata_api() { assert_eq!( @@ -1052,6 +1124,14 @@ mod ddl_tests { assert!(ddl.contains("COMMENT ON COLUMN \"public\".\"users\".\"display_name\" IS 'User''s display name';")); } + + #[test] + fn opengauss_table_ddl_uses_native_tabledef_function() { + assert_eq!( + opengauss_table_ddl_sql("tenant's schema", "active users"), + "SELECT pg_get_tabledef('\"tenant''s schema\".\"active users\"')" + ); + } } pub async fn mysql_ddl(pool: &db::mysql::MySqlPool, table: &str) -> Result { @@ -1086,6 +1166,15 @@ pub async fn sqlite_ddl(pool: &db::sqlite::SqliteHandle, table: &str) -> Result< .map_err(|e| e.to_string())? } +pub async fn opengauss_table_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { + first_string_cell(db::postgres::execute_query(pool, &opengauss_table_ddl_sql(schema, table)).await?) +} + +pub fn opengauss_table_ddl_sql(schema: &str, table: &str) -> String { + let qualified_name = format!("{}.{}", pg_ident(schema), pg_ident(table)); + format!("SELECT pg_get_tabledef({})", sql_string(&qualified_name)) +} + pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { let (columns, indexes, fkeys) = tokio::try_join!( db::postgres::get_columns(pool, schema, table),