From eccbc5a315c1ebe7c973cb11dd00b210c92aa657 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 25 Aug 2026 04:14:09 +0000 Subject: [PATCH] fix(editor): support Spark datasource DDL for Kyuubi Closes #7120 --- .../__tests__/sql/semantic/dialect.spec.ts | 24 +- apps/desktop/src/lib/sql/semantic/dialect.ts | 1 + crates/dbx-core/src/sql_analysis.rs | 252 +++++++++++++++++- crates/dbx-core/tests/sql_analysis.rs | 103 +++++++ 4 files changed, 375 insertions(+), 5 deletions(-) diff --git a/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts b/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts index 491bb3cd08..761c57edad 100644 --- a/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/semantic/dialect.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { sqlSemanticDialectFor } from "@/lib/sql/semantic/dialect"; +import { sqlReferenceAnalysisDialectFor, sqlSemanticDialectFor } from "@/lib/sql/semantic/dialect"; describe("ClickHouse semantic dialect", () => { it("takes precedence over the legacy MySQL behavior dialect", () => { @@ -17,3 +17,25 @@ describe("ClickHouse semantic dialect", () => { expect(dialect.quoteIdentifier("event`name")).toBe("`event``name`"); }); }); + +describe("SQL reference analysis dialect", () => { + it("uses Spark SQL grammar for Kyuubi connections", () => { + expect( + sqlReferenceAnalysisDialectFor({ + databaseType: "kyuubi", + identifierQuote: "`", + fallbackDialect: "generic", + }), + ).toBe("spark"); + }); + + it("preserves the configured fallback for other connection types", () => { + expect( + sqlReferenceAnalysisDialectFor({ + databaseType: "hive", + identifierQuote: "`", + fallbackDialect: "generic", + }), + ).toBe("generic"); + }); +}); diff --git a/apps/desktop/src/lib/sql/semantic/dialect.ts b/apps/desktop/src/lib/sql/semantic/dialect.ts index a62611bee1..67b0e96995 100644 --- a/apps/desktop/src/lib/sql/semantic/dialect.ts +++ b/apps/desktop/src/lib/sql/semantic/dialect.ts @@ -154,6 +154,7 @@ export const SQL_SEMANTIC_DIALECTS: Record = }; export function sqlReferenceAnalysisDialectFor(options: { databaseType?: DatabaseType; identifierQuote?: string; fallbackDialect: string }): string { + if (options.databaseType === "kyuubi") return "spark"; if (options.databaseType === "kingbase" && options.identifierQuote === "`") return "mysql"; return options.fallbackDialect; } diff --git a/crates/dbx-core/src/sql_analysis.rs b/crates/dbx-core/src/sql_analysis.rs index dda2fdf475..112a48f2a3 100644 --- a/crates/dbx-core/src/sql_analysis.rs +++ b/crates/dbx-core/src/sql_analysis.rs @@ -10,6 +10,7 @@ use sqlparser::ast::{ }; use sqlparser::dialect::{ ClickHouseDialect, DuckDbDialect, GenericDialect, MsSqlDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect, + SparkSqlDialect, }; use sqlparser::keywords::Keyword; use sqlparser::parser::{Parser, ParserError}; @@ -83,6 +84,7 @@ struct Analyzer { cte_scope_stack: Vec>, next_scope_id: usize, is_sqlserver: bool, + is_spark: bool, } pub fn analyze_sql_references(sql: &str, dialect: Option<&str>) -> Result { @@ -106,6 +108,7 @@ pub fn analyze_sql_references(sql: &str, dialect: Option<&str>) -> Result parse_sqlserver(&parser_sql), "clickhouse" => Parser::parse_sql(&ClickHouseDialect {}, &parser_sql), "duckdb" => Parser::parse_sql(&DuckDbDialect {}, &parser_sql), + "spark" => parse_spark_sql(&parser_sql), _ => Parser::parse_sql(&GenericDialect {}, &parser_sql), }; let statements = match parsed_statements { @@ -123,7 +126,11 @@ pub fn analyze_sql_references(sql: &str, dialect: Option<&str>) -> Result return Err(error.to_string()), }; - let mut analyzer = Analyzer { is_sqlserver: normalized_dialect == "sqlserver", ..Analyzer::default() }; + let mut analyzer = Analyzer { + is_sqlserver: normalized_dialect == "sqlserver", + is_spark: normalized_dialect == "spark", + ..Analyzer::default() + }; for statement in statements { analyzer.visit_statement(&statement); } @@ -155,6 +162,237 @@ fn parse_sqlserver(sql: &str) -> Result, ParserError> { } } +fn parse_spark_sql(sql: &str) -> Result, ParserError> { + let dialect = SparkSqlDialect {}; + let tokens = Tokenizer::new(&dialect, sql).tokenize_with_location()?; + let mut fallback_tokens = tokens.clone(); + let fallback_changed = normalize_spark_datasource_create_table_tokens(&mut fallback_tokens)?; + match Parser::new(&dialect).with_tokens_with_locations(tokens).parse_statements() { + Ok(statements) => Ok(statements), + Err(error) => { + // sqlparser 0.62 accepts Spark's USING clause but expects later datasource clauses in Hive order. + // The fallback removes only clauses parsed independently from the same located token stream. + if fallback_changed { + if let Ok(statements) = + Parser::new(&dialect).with_tokens_with_locations(fallback_tokens).parse_statements() + { + return Ok(statements); + } + } + Err(error) + } + } +} + +fn normalize_spark_datasource_create_table_tokens(tokens: &mut Vec) -> Result { + let dialect = SparkSqlDialect {}; + let mut removals = Vec::new(); + let mut index = 0usize; + + while let Some(statement_start) = next_significant_token_index(tokens, index) { + if matches!(tokens[statement_start].token, Token::EOF) { + break; + } + if matches!(tokens[statement_start].token, Token::SemiColon) { + index = statement_start + 1; + continue; + } + + let statement_end = sql_statement_end_index(tokens, statement_start); + let significant_indexes = tokens + .iter() + .enumerate() + .take(statement_end) + .skip(statement_start) + .filter_map(|(index, token)| (!matches!(token.token, Token::Whitespace(_))).then_some(index)) + .collect::>(); + removals.extend(spark_datasource_clause_removals(tokens, &significant_indexes, &dialect)?); + index = statement_end.saturating_add(1); + } + + removals.sort_unstable(); + removals.dedup(); + let changed = !removals.is_empty(); + for index in removals.into_iter().rev() { + tokens.remove(index); + } + Ok(changed) +} + +fn spark_datasource_clause_removals( + tokens: &[TokenWithSpan], + indexes: &[usize], + dialect: &SparkSqlDialect, +) -> Result, ParserError> { + let Some(using_position) = spark_create_table_using_position(tokens, indexes) else { + return Ok(vec![]); + }; + let Some(provider_index) = indexes.get(using_position + 1).copied() else { + return Ok(vec![]); + }; + if !matches!(tokens[provider_index].token, Token::Word(_)) { + return Parser::new(dialect).expected("data source identifier", tokens[provider_index].clone()); + } + + let mut removals = Vec::new(); + let mut seen_options = false; + let mut seen_partitioning = false; + let mut seen_comment = false; + let mut seen_table_properties = false; + let mut position = using_position + 2; + while position < indexes.len() { + let token = &tokens[indexes[position]]; + match unquoted_token_keyword(token) { + Some(Keyword::OPTIONS) | Some(Keyword::TBLPROPERTIES) => { + let keyword = unquoted_token_keyword(token).expect("matched option keyword"); + let (seen, clause_name) = match keyword { + Keyword::OPTIONS => (&mut seen_options, "OPTIONS"), + Keyword::TBLPROPERTIES => (&mut seen_table_properties, "TBLPROPERTIES"), + _ => unreachable!("matched option keyword"), + }; + if *seen { + return Parser::new(dialect).expected(&format!("at most one {clause_name} clause"), token.clone()); + } + *seen = true; + let Some(open_index) = indexes.get(position + 1).copied() else { + return Ok(vec![]); + }; + if !matches!(tokens[open_index].token, Token::LParen) { + return Ok(vec![]); + } + let Some(close_index) = matching_parenthesis_index(tokens, open_index) else { + return Ok(vec![]); + }; + let Some(close_position) = indexes.iter().position(|index| *index == close_index) else { + return Ok(vec![]); + }; + if !spark_tokens_parse_options(tokens, &indexes[position..=close_position], keyword, dialect) { + return Ok(vec![]); + } + removals.extend_from_slice(&indexes[position..=close_position]); + position = close_position + 1; + } + Some(Keyword::PARTITIONED) => { + if seen_partitioning { + return Parser::new(dialect).expected("at most one PARTITIONED BY clause", token.clone()); + } + seen_partitioning = true; + if indexes.get(position + 1).and_then(|index| unquoted_token_keyword(&tokens[*index])) + != Some(Keyword::BY) + { + return Ok(vec![]); + } + let Some(open_index) = indexes.get(position + 2).copied() else { + return Ok(vec![]); + }; + if !matches!(tokens[open_index].token, Token::LParen) { + return Ok(vec![]); + } + let Some(close_index) = matching_parenthesis_index(tokens, open_index) else { + return Ok(vec![]); + }; + let Some(close_position) = indexes.iter().position(|index| *index == close_index) else { + return Ok(vec![]); + }; + if !spark_tokens_parse_partitioning(tokens, &indexes[position..=close_position], dialect) { + return Ok(vec![]); + } + removals.extend_from_slice(&indexes[position..=close_position]); + position = close_position + 1; + } + Some(Keyword::COMMENT) => { + if seen_comment { + return Parser::new(dialect).expected("at most one COMMENT clause", token.clone()); + } + seen_comment = true; + let Some(comment_index) = indexes.get(position + 1).copied() else { + return Ok(vec![]); + }; + if !matches!(tokens[comment_index].token, Token::SingleQuotedString(_)) { + return Ok(vec![]); + } + removals.extend_from_slice(&indexes[position..=position + 1]); + position += 2; + } + Some(Keyword::LOCATION) => { + let Some(location_index) = indexes.get(position + 1).copied() else { + return Ok(vec![]); + }; + if !matches!(tokens[location_index].token, Token::SingleQuotedString(_)) { + return Ok(vec![]); + } + position += 2; + } + Some(Keyword::AS) => break, + _ => return Ok(vec![]), + } + } + Ok(removals) +} + +fn spark_create_table_using_position(tokens: &[TokenWithSpan], indexes: &[usize]) -> Option { + if indexes.first().and_then(|index| unquoted_token_keyword(&tokens[*index])) != Some(Keyword::CREATE) { + return None; + } + let mut position = 1usize; + if indexes.get(position).and_then(|index| unquoted_token_keyword(&tokens[*index])) == Some(Keyword::OR) + && indexes.get(position + 1).and_then(|index| unquoted_token_keyword(&tokens[*index])) == Some(Keyword::REPLACE) + { + position += 2; + } + while indexes.get(position).is_some_and(|index| { + ["GLOBAL", "LOCAL", "TEMP", "TEMPORARY", "EXTERNAL"] + .iter() + .any(|modifier| unquoted_token_word_eq(&tokens[*index], modifier)) + }) { + position += 1; + } + if indexes.get(position).and_then(|index| unquoted_token_keyword(&tokens[*index])) != Some(Keyword::TABLE) { + return None; + } + + let mut depth = 0usize; + for (position, index) in indexes.iter().copied().enumerate().skip(position + 1) { + match tokens[index].token { + Token::LParen => depth += 1, + Token::RParen => depth = depth.saturating_sub(1), + _ if depth == 0 && unquoted_token_keyword(&tokens[index]) == Some(Keyword::AS) => return None, + _ if depth == 0 && unquoted_token_keyword(&tokens[index]) == Some(Keyword::USING) => { + return Some(position); + } + _ => {} + } + } + None +} + +fn spark_tokens_parse_options( + tokens: &[TokenWithSpan], + indexes: &[usize], + keyword: Keyword, + dialect: &SparkSqlDialect, +) -> bool { + let clause_tokens = indexes.iter().map(|index| tokens[*index].clone()).collect(); + let mut parser = Parser::new(dialect).with_tokens_with_locations(clause_tokens); + parser.parse_options(keyword).is_ok() && matches!(parser.peek_token_ref().token, Token::EOF) +} + +fn spark_tokens_parse_partitioning(tokens: &[TokenWithSpan], indexes: &[usize], dialect: &SparkSqlDialect) -> bool { + let clause_tokens = indexes.iter().map(|index| tokens[*index].clone()).collect(); + let mut parser = Parser::new(dialect).with_tokens_with_locations(clause_tokens); + if parser.expect_keywords(&[Keyword::PARTITIONED, Keyword::BY]).is_err() + || parser.expect_token(&Token::LParen).is_err() + { + return false; + } + let Ok(expressions) = parser.parse_comma_separated(Parser::parse_expr) else { + return false; + }; + !expressions.is_empty() + && parser.expect_token(&Token::RParen).is_ok() + && matches!(parser.peek_token_ref().token, Token::EOF) +} + fn normalize_sqlserver_cursor_tokens(tokens: &mut Vec) -> bool { let mut removals = Vec::new(); let mut index = 0usize; @@ -168,7 +406,7 @@ fn normalize_sqlserver_cursor_tokens(tokens: &mut Vec) -> bool { continue; } - let statement_end = sqlserver_statement_end_index(tokens, statement_start); + let statement_end = sql_statement_end_index(tokens, statement_start); let significant_indexes = tokens .iter() .enumerate() @@ -275,7 +513,7 @@ fn normalize_sqlserver_alter_table_add_tokens(tokens: &mut Vec) - continue; } - let statement_end = sqlserver_statement_end_index(tokens, table_index + 1); + let statement_end = sql_statement_end_index(tokens, table_index + 1); let Some(add_index) = sqlserver_alter_table_add_index(tokens, table_index + 1, statement_end) else { index = statement_end.saturating_add(1); continue; @@ -311,7 +549,7 @@ fn normalize_sqlserver_alter_table_add_tokens(tokens: &mut Vec) - changed } -fn sqlserver_statement_end_index(tokens: &[TokenWithSpan], start: usize) -> usize { +fn sql_statement_end_index(tokens: &[TokenWithSpan], start: usize) -> usize { let mut depth = 0usize; for (index, token) in tokens.iter().enumerate().skip(start) { match token.token { @@ -771,6 +1009,7 @@ fn normalize_dialect(dialect: Option<&str>) -> String { "sqlserver" | "mssql" => "sqlserver".to_string(), "clickhouse" => "clickhouse".to_string(), "duckdb" => "duckdb".to_string(), + "spark" | "sparksql" => "spark".to_string(), _ => "generic".to_string(), } } @@ -779,6 +1018,11 @@ impl Analyzer { fn visit_statement(&mut self, statement: &Statement) { match statement { Statement::Query(query) => self.visit_query_in_new_scope(query, None), + Statement::CreateTable(create_table) if self.is_spark => { + if let Some(query) = &create_table.query { + self.visit_query_in_new_scope(query, None); + } + } Statement::Declare { stmts } if self.is_sqlserver => { for declaration in stmts { if let Some(query) = &declaration.for_query { diff --git a/crates/dbx-core/tests/sql_analysis.rs b/crates/dbx-core/tests/sql_analysis.rs index 8d8ba21c3a..b03d8f05b0 100644 --- a/crates/dbx-core/tests/sql_analysis.rs +++ b/crates/dbx-core/tests/sql_analysis.rs @@ -696,3 +696,106 @@ fn clickhouse_strictness_first_left_joins_do_not_raise_syntax_errors() { assert_eq!(tables, vec![("events", Some("a")), ("wallets", Some("b"))]); } } + +#[test] +fn spark_datasource_create_table_supports_iceberg_clauses() { + let sql = r#"CREATE TABLE account_flow ( + id STRING, + databasename STRING, + created TIMESTAMP +) +USING iceberg +PARTITIONED BY (databasename, truncate(created, 7)) +COMMENT '账户流水表' +TBLPROPERTIES ( + 'format-version' = '2', + 'snapshot.base.keep.minutes' = '1440', + 'self-optimizing.group' = 'supbig', + 'write.metadata.delete-after-commit.enabled' = 'true', + 'write.metadata.previous-versions-max' = '3', + 'clean-orphan-file.enabled' = 'true', + 'clean-orphan-file.min-existing-time-minutes' = '1440', + 'primary-key' = 'id,databasename', + 'table.drop-base-path.enabled' = 'true' +);"#; + + let analysis = analyze_sql_references(sql, Some("spark")) + .unwrap_or_else(|error| panic!("Spark datasource CREATE TABLE should analyze: {error}")); + + assert!(analysis.tables.is_empty()); + assert!(analysis.columns.is_empty()); +} + +#[test] +fn spark_datasource_ctas_preserves_query_references() { + let sql = "CREATE TABLE account_flow USING iceberg PARTITIONED BY (id) TBLPROPERTIES ('format-version' = '2') AS SELECT s.id FROM source_flow s"; + let analysis = analyze_sql_references(sql, Some("spark")).expect("Spark datasource CTAS should analyze"); + + let tables: Vec<_> = analysis.tables.iter().map(|table| (table.name.as_str(), table.alias.as_deref())).collect(); + assert_eq!(tables, vec![("source_flow", Some("s"))]); + assert_eq!(analysis.tables[0].span.start_column, sql.find("source_flow").expect("source table") + 1); + assert_eq!(analysis.columns.iter().map(|column| column.name.as_str()).collect::>(), vec!["id"]); +} + +#[test] +fn spark_datasource_create_table_validates_options_and_properties() { + let sql = "CREATE TABLE account_flow (id STRING) USING iceberg TBLPROPERTIES ('format-version' = '2') COMMENT 'account flow' PARTITIONED BY (id) OPTIONS ('merge-schema' = 'true')"; + let analysis = + analyze_sql_references(sql, Some("spark")).expect("reordered Spark datasource clauses should analyze"); + + assert!(analysis.tables.is_empty()); + assert!(analysis.columns.is_empty()); +} + +#[test] +fn spark_datasource_create_table_rejects_duplicate_clauses() { + let duplicate_sql = [ + "CREATE TABLE broken (id STRING) USING iceberg OPTIONS ('a' = '1') OPTIONS ('b' = '2')", + "CREATE TABLE broken (id STRING) USING iceberg PARTITIONED BY (id) PARTITIONED BY (id)", + "CREATE TABLE broken (id STRING) USING iceberg COMMENT 'first' COMMENT 'second'", + "CREATE TABLE broken (id STRING) USING iceberg TBLPROPERTIES ('a' = '1') TBLPROPERTIES ('b' = '2')", + ]; + + let accepted: Vec<_> = + duplicate_sql.iter().copied().filter(|sql| analyze_sql_references(sql, Some("spark")).is_ok()).collect(); + assert!(accepted.is_empty(), "duplicate Spark datasource clauses were accepted: {accepted:?}"); +} + +#[test] +fn spark_datasource_create_table_keeps_syntax_errors() { + for sql in [ + "CREATE TABLE broken (id STRING) USING", + "CREATE TABLE broken (id STRING) USING 'iceberg'", + "CREATE TABLE broken (id STRING USING iceberg", + "CREATE TABLE broken (id STRING, created TIMESTAMP) USING iceberg PARTITIONED BY (truncate(created 7))", + "CREATE TABLE broken (id STRING, created TIMESTAMP) USING iceberg PARTITIONED BY (created) UNKNOWN CLAUSE", + "CREATE TABLE broken (id STRING) USING iceberg OPTIONS ('merge-schema')", + "CREATE TABLE broken (id STRING) USING iceberg TBLPROPERTIES ('format-version')", + ] { + let error = analyze_sql_references(sql, Some("spark")) + .expect_err(&format!("malformed Spark datasource CREATE TABLE must keep its parser error: {sql}")); + assert!(!error.is_empty()); + } + + let provider_error = analyze_sql_references("CREATE TABLE broken (id STRING) USING 'iceberg'", Some("spark")) + .expect_err("quoted Spark datasource provider must remain invalid"); + assert!(provider_error.contains("Line: 1, Column:")); +} + +#[test] +fn spark_selects_still_report_query_references() { + let analysis = analyze_sql_references("SELECT s.id FROM source_flow s", Some("spark")) + .expect("ordinary Spark SELECT should analyze"); + + assert_eq!(analysis.tables[0].name, "source_flow"); + assert_eq!(analysis.tables[0].alias.as_deref(), Some("s")); + assert_eq!(analysis.columns[0].name, "id"); +} + +#[test] +fn generic_dialect_still_rejects_spark_datasource_clauses() { + let error = analyze_sql_references("CREATE TABLE account_flow (id STRING) USING iceberg", Some("generic")) + .expect_err("Spark datasource clauses must remain dialect-specific"); + + assert!(error.contains("USING")); +}