From f83f707fbcd0466bdec254de7ec9c1d9845c15cf Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 28 Aug 2026 09:13:09 +0000 Subject: [PATCH] fix(kingbase): honor driver identifier quote for schema context Closes #7505 --- .../drivers/kingbase-go/kingbase_metadata.go | 5 +- agents/drivers/kingbase-go/main.go | 9 +- agents/drivers/kingbase-go/main_test.go | 33 ++++++ crates/dbx-core/src/connection.rs | 10 +- crates/dbx-core/src/db/agent_driver.rs | 19 +++- crates/dbx-core/src/query.rs | 102 +++++++++++++++--- 6 files changed, 159 insertions(+), 19 deletions(-) diff --git a/agents/drivers/kingbase-go/kingbase_metadata.go b/agents/drivers/kingbase-go/kingbase_metadata.go index c3b04326f2..78df2a319a 100644 --- a/agents/drivers/kingbase-go/kingbase_metadata.go +++ b/agents/drivers/kingbase-go/kingbase_metadata.go @@ -222,10 +222,7 @@ func (s *server) identifierQuote() string { // schema/table names containing hyphens or other special characters render as // valid SQL instead of being parsed as operators or bare tokens. func (s *server) quoteDDLIdentifier(value string) string { - if s.mode.mysqlCompat { - return "`" + strings.ReplaceAll(value, "`", "``") + "`" - } - return quoteIdentifier(value) + return s.quoteIdentifier(value) } func (s *server) connectionInfo() (map[string]any, error) { diff --git a/agents/drivers/kingbase-go/main.go b/agents/drivers/kingbase-go/main.go index b66a3f15d5..a9cfae19c4 100644 --- a/agents/drivers/kingbase-go/main.go +++ b/agents/drivers/kingbase-go/main.go @@ -934,7 +934,7 @@ func (s *server) setSchema(ctx context.Context, conn *sql.Conn, schema string) e if schema != "" { // Kingbase implicitly prioritizes its system catalog when it is not // listed explicitly, matching the JDBC agent and DBeaver behavior. - statement = "SET search_path TO " + quoteIdentifier(schema) + statement = "SET search_path TO " + s.quoteIdentifier(schema) } if _, err := conn.ExecContext(ctx, statement); err != nil { return err @@ -1893,6 +1893,13 @@ func quoteIdentifier(value string) string { return `"` + strings.ReplaceAll(value, `"`, `""`) + `"` } +func (s *server) quoteIdentifier(value string) string { + if s.mode.mysqlCompat { + return "`" + strings.ReplaceAll(value, "`", "``") + "`" + } + return quoteIdentifier(value) +} + func quoteLiteral(value string) string { return "'" + strings.ReplaceAll(value, "'", "''") + "'" } diff --git a/agents/drivers/kingbase-go/main_test.go b/agents/drivers/kingbase-go/main_test.go index dc0d892253..5e53345ee4 100644 --- a/agents/drivers/kingbase-go/main_test.go +++ b/agents/drivers/kingbase-go/main_test.go @@ -3954,6 +3954,39 @@ func TestSchemaConnResetsOnceAfterExplicitSchema(t *testing.T) { } } +func TestSchemaConnUsesBackticksInMySQLCompatMode(t *testing.T) { + db, state := openFakeDB(t, 0) + server := newServer() + server.db = db + server.mode.mysqlCompat = true + + conn, err := server.schemaConn(context.Background(), "audit-schema") + if err != nil { + t.Fatal(err) + } + if err := conn.Close(); err != nil { + t.Fatal(err) + } + + state.mu.Lock() + defer state.mu.Unlock() + if len(state.execStatements) != 1 || state.execStatements[0] != "SET search_path TO `audit-schema`" { + t.Fatalf("unexpected MySQL compatibility schema setup: %v", state.execStatements) + } +} + +func TestKingbaseIdentifierQuoteEscapesModeSpecificDelimiter(t *testing.T) { + server := newServer() + server.mode.mysqlCompat = true + if got := server.quoteIdentifier("audit`schema"); got != "`audit``schema`" { + t.Fatalf("unexpected MySQL compatibility identifier: %s", got) + } + server.mode.mysqlCompat = false + if got := server.quoteIdentifier(`audit"schema`); got != `"audit""schema"` { + t.Fatalf("unexpected PostgreSQL-compatible identifier: %s", got) + } +} + func TestSchemaConnPropagatesSchemaErrors(t *testing.T) { tests := []struct { name string diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index d7fb973321..a25a665d6e 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2455,7 +2455,7 @@ impl AppState { } } } - let client = match initial_result { + let mut client = match initial_result { Ok(client) => client, Err(err) => { let alternate_configs = oracle_alternate_connect_configs(&db_config, &err); @@ -2529,6 +2529,14 @@ impl AppState { } } }; + if db_config.db_type == DatabaseType::Kingbase { + let identifier_quote = client + .connection_info(Some(agent_connect_timeout(&db_config))) + .await + .ok() + .map(|info| info.identifier_quote); + client.set_identifier_quote(identifier_quote); + } PoolKind::agent(client) } else { // ZooKeeper JVM properties are connection-scoped; shared agent daemons must not inherit them. diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index b8deae1462..3d1260ed4c 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -887,6 +887,9 @@ pub struct AgentDriverClient { shared_runtime: Option>, agent_session_id: Option, cached_query: Option, + /// Driver-reported identifier quoting, captured once after connect. + /// Keeping this on the client avoids a connection-info RPC on every query. + identifier_quote: Option, } /// Keeps serialized session RPC access separate from the process-level fail-stop handle. @@ -896,13 +899,19 @@ pub struct PooledAgentClient { client: tokio::sync::Mutex, shared_runtime: Option>, agent_session_id: Option, + identifier_quote: Option, } impl PooledAgentClient { pub fn new(client: AgentDriverClient) -> Self { let shared_runtime = client.shared_runtime.clone(); let agent_session_id = client.agent_session_id.clone(); - Self { client: tokio::sync::Mutex::new(client), shared_runtime, agent_session_id } + let identifier_quote = client.identifier_quote.clone(); + Self { client: tokio::sync::Mutex::new(client), shared_runtime, agent_session_id, identifier_quote } + } + + pub fn identifier_quote(&self) -> Option<&str> { + self.identifier_quote.as_deref() } pub async fn lock(&self) -> tokio::sync::MutexGuard<'_, AgentDriverClient> { @@ -1632,6 +1641,7 @@ impl AgentDriverClient { shared_runtime: None, agent_session_id: None, cached_query: None, + identifier_quote: None, }) } @@ -1646,6 +1656,7 @@ impl AgentDriverClient { shared_runtime: Some(runtime), agent_session_id: Some(agent_session_id), cached_query: None, + identifier_quote: None, } } @@ -1947,6 +1958,10 @@ impl AgentDriverClient { self.call_method_with_timeout(AgentMethod::ConnectionInfo, serde_json::json!({}), timeout_duration).await } + pub(crate) fn set_identifier_quote(&mut self, identifier_quote: Option) { + self.identifier_quote = identifier_quote.filter(|quote| !quote.trim().is_empty()); + } + pub async fn disconnect(&mut self) -> Result { self.invalidate_cached_query(); if self.shared_runtime.is_some() { @@ -3345,6 +3360,7 @@ impl AgentDriverClient { shared_runtime: None, agent_session_id: None, cached_query: None, + identifier_quote: None, } } @@ -4070,6 +4086,7 @@ mod tests { shared_runtime: None, agent_session_id: None, cached_query: None, + identifier_quote: None, }; let started_at = std::time::Instant::now(); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 6348ccb67c..80a3111925 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -826,13 +826,22 @@ fn schema_for_execution_context(db_type: Option, schema: Option<&s } fn sql_for_execution_context(db_type: Option, sql: &str, schema: Option<&str>) -> String { + sql_for_execution_context_with_identifier_quote(db_type, sql, schema, None) +} + +fn sql_for_execution_context_with_identifier_quote( + db_type: Option, + sql: &str, + schema: Option<&str>, + identifier_quote: Option<&str>, +) -> String { let Some(schema) = schema.map(str::trim).filter(|schema| !schema.is_empty()) else { return sql.to_string(); }; match db_type { Some(DatabaseType::Iris) => qualify_iris_unqualified_dml(sql, schema).unwrap_or_else(|| sql.to_string()), Some(DatabaseType::Kingbase) => { - qualify_kingbase_unqualified_relations(sql, schema).unwrap_or_else(|| sql.to_string()) + qualify_kingbase_unqualified_relations(sql, schema, identifier_quote).unwrap_or_else(|| sql.to_string()) } _ => sql.to_string(), } @@ -862,7 +871,7 @@ fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option { changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) } -fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str) -> Option { +fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str, identifier_quote: Option<&str>) -> Option { let dialect = PostgreSqlDialect {}; let mut statements = Parser::parse_sql(&dialect, sql).ok()?; if statements.is_empty() { @@ -875,8 +884,13 @@ fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str) -> Option Option { schema: &'a str, + identifier_quote: char, cte_names: &'a HashSet, parameterized_table_depth: usize, changed: bool, @@ -910,7 +925,12 @@ impl VisitorMut for KingbaseRelationQualifier<'_> { fn post_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow { if self.parameterized_table_depth == 0 - && qualify_unqualified_relation_name(relation, self.schema, self.cte_names) + && qualify_unqualified_relation_name_with_quote( + relation, + self.schema, + self.cte_names, + self.identifier_quote, + ) { self.changed = true; } @@ -930,6 +950,15 @@ fn statement_uses_schema_context(statement: &Statement) -> bool { } fn qualify_unqualified_relation_name(name: &mut ObjectName, schema: &str, cte_names: &HashSet) -> bool { + qualify_unqualified_relation_name_with_quote(name, schema, cte_names, '"') +} + +fn qualify_unqualified_relation_name_with_quote( + name: &mut ObjectName, + schema: &str, + cte_names: &HashSet, + identifier_quote: char, +) -> bool { let [ObjectNamePart::Identifier(table)] = name.0.as_slice() else { return false; }; @@ -938,10 +967,21 @@ fn qualify_unqualified_relation_name(name: &mut ObjectName, schema: &str, cte_na } let table = table.clone(); - name.0 = vec![ObjectNamePart::Identifier(Ident::with_quote('"', schema)), ObjectNamePart::Identifier(table)]; + name.0 = vec![ + ObjectNamePart::Identifier(Ident::with_quote(identifier_quote, schema)), + ObjectNamePart::Identifier(table), + ]; true } +fn identifier_quote_char(identifier_quote: Option<&str>) -> char { + match identifier_quote.map(str::trim) { + Some("`") => '`', + Some("\"") => '"', + _ => '"', + } +} + fn statement_cte_names(statement: &Statement) -> HashSet { let mut names = HashSet::new(); collect_statement_cte_names(statement, &mut names); @@ -2055,7 +2095,8 @@ async fn do_execute_typed( PoolKind::Agent(client) => { let client = client.clone(); let source_client = client.clone(); - let sql = sql_for_execution_context(pool_db_type, sql, schema); + let sql = + sql_for_execution_context_with_identifier_quote(pool_db_type, sql, schema, client.identifier_quote()); let database = database.map(|s| s.to_string()); let schema = schema_for_execution_context(pool_db_type, schema).map(|s| s.to_string()); let max_rows = options.max_rows; @@ -3616,8 +3657,12 @@ pub async fn execute_statements( let execution_schema = schema_for_execution_context(db_type, schema); let rewritten_statements; let statements = if qualifies_unqualified_agent_relations(db_type) { - rewritten_statements = - statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::>(); + rewritten_statements = statements + .iter() + .map(|sql| { + sql_for_execution_context_with_identifier_quote(db_type, sql, schema, client.identifier_quote()) + }) + .collect::>(); rewritten_statements.as_slice() } else { statements @@ -4492,8 +4537,10 @@ async fn exec_tx_agent_inner( let execution_schema = schema_for_execution_context(db_type, schema); let rewritten_statements; let statements = if qualifies_unqualified_agent_relations(db_type) { - rewritten_statements = - statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::>(); + rewritten_statements = statements + .iter() + .map(|sql| sql_for_execution_context_with_identifier_quote(db_type, sql, schema, client.identifier_quote())) + .collect::>(); rewritten_statements.as_slice() } else { statements @@ -5364,7 +5411,7 @@ async fn execute_manual_txn_agent_statement( page_size: Option, result_session_id: Option<&str>, ) -> Result { - let sql = sql_for_execution_context(db_type, statement, schema); + let sql = sql_for_execution_context_with_identifier_quote(db_type, statement, schema, client.identifier_quote()); let execution_schema = schema_for_execution_context(db_type, schema); let options = manual_txn_agent_query_options(row_limit, table_data_preview, page_size, result_session_id); let request = manual_txn_agent_query_request( @@ -8381,6 +8428,37 @@ for line in sys.stdin: ); } + #[test] + fn kingbase_execution_context_uses_driver_identifier_quote() { + assert_eq!( + sql_for_execution_context_with_identifier_quote( + Some(DatabaseType::Kingbase), + "SELECT * FROM sys_user", + Some("audit-schema"), + Some("`"), + ), + "SELECT * FROM `audit-schema`.sys_user" + ); + assert_eq!( + sql_for_execution_context_with_identifier_quote( + Some(DatabaseType::Kingbase), + "SELECT * FROM sys_user", + Some("audit-schema"), + Some("\""), + ), + "SELECT * FROM \"audit-schema\".sys_user" + ); + assert_eq!( + sql_for_execution_context_with_identifier_quote( + Some(DatabaseType::Kingbase), + "SELECT * FROM sys_user", + Some("audit-schema"), + Some("unsupported"), + ), + "SELECT * FROM \"audit-schema\".sys_user" + ); + } + #[test] fn parses_postgres_drop_database_target() { assert_eq!(parse_drop_database_target("DROP DATABASE vaultwarden;"), Some("vaultwarden".to_string()));