feat(query): move SQL editability analysis to Rust

This commit is contained in:
t8y2 2026-05-22 15:07:26 +08:00
parent c736fe372a
commit cb3e1d6e9f
12 changed files with 685 additions and 7 deletions

View File

@ -87,6 +87,7 @@ export const executeInTransaction = forward("executeInTransaction");
export const cancelQuery = forward("cancelQuery");
export const closeQuerySession = forward("closeQuerySession");
export const analyzeSqlReferences = forward("analyzeSqlReferences");
export const analyzeEditableQueryEditability = forward("analyzeEditableQueryEditability");
// AI
export const aiComplete = forward("aiComplete");

View File

@ -49,6 +49,7 @@ import type {
ExportProgress,
XlsxCellValue,
} from "./tauri";
import type { QueryEditability } from "@/lib/sqlAnalysis";
// ---------------------------------------------------------------------------
// Helpers
@ -404,6 +405,10 @@ export async function analyzeSqlReferences(sql: string, dialect?: string): Promi
return post("/api/query/analyze-sql-references", { sql, dialect });
}
export async function analyzeEditableQueryEditability(sql: string): Promise<QueryEditability> {
return post("/api/query/analyze-editability", { sql });
}
// ---------------------------------------------------------------------------
// AI
// ---------------------------------------------------------------------------

View File

@ -21,6 +21,7 @@ import type {
SavedSqlLibrary,
} from "@/types/database";
import type { AiConfig } from "@/stores/settingsStore";
import type { QueryEditability } from "@/lib/sqlAnalysis";
export interface AgentDriverInfo {
db_type: string;
@ -312,6 +313,10 @@ export async function analyzeSqlReferences(sql: string, dialect?: string): Promi
return invoke("analyze_sql_references", { sql, dialect });
}
export async function analyzeEditableQueryEditability(sql: string): Promise<QueryEditability> {
return invoke("analyze_editable_query_editability", { sql });
}
export async function listIndexes(
connectionId: string,
database: string,

View File

@ -6,12 +6,7 @@ import { orderPinnedFirst } from "@/lib/pinnedItems";
import { canCancelQueryExecution } from "@/lib/queryExecutionState";
import { closeAllTabsState, closeOtherTabsState } from "@/lib/tabCloseActions";
import { buildExplainSql, parseExplainResult } from "@/lib/explainPlan";
import {
allEditableColumnsWriteable,
allPrimaryKeysPresent,
analyzeEditableQueryEditability,
sourceColumnsForResult,
} from "@/lib/sqlAnalysis";
import { allEditableColumnsWriteable, allPrimaryKeysPresent, sourceColumnsForResult } from "@/lib/sqlAnalysis";
import { restoreOpenTabsState, serializeOpenTabs } from "@/lib/openTabsPersistence";
import {
mongoCountToQueryResult,
@ -373,7 +368,7 @@ export const useQueryStore = defineStore("query", () => {
return;
}
const editability = analyzeEditableQueryEditability(sql);
const editability = await api.analyzeEditableQueryEditability(sql);
if (!editability.editable) {
tab.queryAnalysis = undefined;
tab.querySourceColumns = undefined;

View File

@ -19,6 +19,7 @@ pub mod saved_sql;
pub mod schema;
pub mod sql;
pub mod sql_analysis;
pub mod sql_editability;
pub mod storage;
pub mod table_import;
pub mod transfer;

View File

@ -0,0 +1,616 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct EditableQueryInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub schema: Option<String>,
pub table_name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub table_alias: Option<String>,
pub select_star: bool,
pub columns: Vec<EditableQueryColumn>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct EditableQueryColumn {
#[serde(skip_serializing_if = "Option::is_none")]
pub source_name: Option<String>,
pub result_name: String,
pub expression: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum QueryEditabilityReason {
NotSelect,
Cte,
SetOperation,
Aggregation,
ComplexSource,
ComputedColumns,
NoTable,
NoPrimaryKey,
PrimaryKeyNotReturned,
AliasedColumns,
MetadataUnavailable,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct QueryEditability {
pub editable: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub analysis: Option<EditableQueryInfo>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reason: Option<QueryEditabilityReason>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct FromSource {
schema: Option<String>,
table_name: String,
alias: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct QualifiedIdentifier {
parts: Vec<String>,
end: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct Identifier {
value: String,
end: usize,
}
pub fn analyze_editable_query(sql: &str) -> Option<EditableQueryInfo> {
let result = analyze_editable_query_editability(sql);
if result.editable {
result.analysis
} else {
None
}
}
pub fn analyze_editable_query_editability(sql: &str) -> QueryEditability {
let normalized = strip_sql_comments(sql).trim_end_matches(';').trim().to_string();
if normalized.is_empty() {
return not_editable(QueryEditabilityReason::NotSelect);
}
if starts_with_keyword(&normalized, "WITH") {
return not_editable(QueryEditabilityReason::Cte);
}
if !starts_with_keyword(&normalized, "SELECT") {
return not_editable(QueryEditabilityReason::NotSelect);
}
if has_top_level_keyword(&normalized, &["UNION", "INTERSECT", "EXCEPT"]) {
return not_editable(QueryEditabilityReason::SetOperation);
}
if normalized.contains(';') {
return not_editable(QueryEditabilityReason::ComplexSource);
}
let Some(from_index) = find_top_level_keyword(&normalized, "FROM", 0) else {
return not_editable(QueryEditabilityReason::NoTable);
};
let select_body = normalized["SELECT".len()..from_index].trim();
if starts_with_keyword(select_body, "DISTINCT") {
return not_editable(QueryEditabilityReason::Aggregation);
}
let group_index = find_top_level_keyword(&normalized, "GROUP", from_index + "FROM".len());
let having_index = find_top_level_keyword(&normalized, "HAVING", from_index + "FROM".len());
if group_index.is_some() || having_index.is_some() {
return not_editable(QueryEditabilityReason::Aggregation);
}
let from_body_start = from_index + "FROM".len();
let from_end =
first_top_level_keyword_index(&normalized, &["WHERE", "ORDER", "LIMIT", "OFFSET", "FETCH"], from_body_start)
.unwrap_or(normalized.len());
let from_body = normalized[from_body_start..from_end].trim();
let Some(source) = parse_from_source(from_body) else {
return not_editable(QueryEditabilityReason::ComplexSource);
};
let select_star = is_select_star(select_body, source.alias.as_deref());
let columns = if select_star { Vec::new() } else { parse_select_columns(select_body) };
if !select_star && columns.is_empty() {
return not_editable(QueryEditabilityReason::ComputedColumns);
}
QueryEditability {
editable: true,
analysis: Some(EditableQueryInfo {
schema: source.schema,
table_name: source.table_name,
table_alias: source.alias,
select_star,
columns,
}),
reason: None,
}
}
fn not_editable(reason: QueryEditabilityReason) -> QueryEditability {
QueryEditability { editable: false, analysis: None, reason: Some(reason) }
}
fn parse_select_columns(body: &str) -> Vec<EditableQueryColumn> {
let mut columns = Vec::new();
let mut depth = 0i32;
let mut current = String::new();
let mut quote: Option<char> = None;
for ch in body.chars() {
if let Some(close) = quote {
current.push(ch);
if ch == close {
quote = None;
}
continue;
}
match ch {
'\'' | '"' | '`' => quote = Some(ch),
'[' => quote = Some(']'),
'(' => depth += 1,
')' => depth -= 1,
',' if depth == 0 => {
let Some(column) = parse_select_column(current.trim()) else {
return Vec::new();
};
columns.push(column);
current.clear();
continue;
}
_ => {}
}
current.push(ch);
}
if !current.trim().is_empty() {
let Some(column) = parse_select_column(current.trim()) else {
return Vec::new();
};
columns.push(column);
}
columns
}
fn parse_select_column(column: &str) -> Option<EditableQueryColumn> {
let Some(source) = parse_qualified_identifier(column) else {
return parse_computed_select_column(column);
};
let rest = &column[source.end..];
let Some(alias) = parse_column_alias(rest) else {
return parse_computed_select_column(column);
};
let source_name = source.parts.last()?.clone();
Some(EditableQueryColumn {
source_name: Some(source_name.clone()),
result_name: alias.unwrap_or(source_name),
expression: column[..source.end].trim().to_string(),
})
}
fn parse_computed_select_column(column: &str) -> Option<EditableQueryColumn> {
let alias = parse_expression_alias(column)?;
Some(EditableQueryColumn { source_name: None, result_name: alias.result_name, expression: alias.expression })
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ExpressionAlias {
expression: String,
result_name: String,
}
fn parse_expression_alias(column: &str) -> Option<ExpressionAlias> {
let trimmed_end = column.trim_end();
for (index, _) in trimmed_end.match_indices(|c: char| c == 'A' || c == 'a') {
let candidate = &trimmed_end[index..];
if !candidate.get(..2).is_some_and(|prefix| prefix.eq_ignore_ascii_case("AS")) {
continue;
}
let before = if index == 0 { "" } else { &trimmed_end[..index] };
if before.chars().last().is_some_and(is_identifier_char) {
continue;
}
let after_as = &candidate[2..];
if !after_as.chars().next().is_some_and(char::is_whitespace) {
continue;
}
let alias_text = after_as.trim();
let alias = read_identifier(alias_text, 0)?;
if alias.end != alias_text.len() {
continue;
}
let expression = trimmed_end[..index].trim().to_string();
if expression.is_empty() {
return None;
}
return Some(ExpressionAlias { expression, result_name: alias.value });
}
None
}
fn parse_column_alias(rest: &str) -> Option<Option<String>> {
let trimmed = rest.trim();
if trimmed.is_empty() {
return Some(None);
}
let alias_text = strip_leading_as(trimmed).unwrap_or(trimmed).trim();
let alias = read_identifier(alias_text, 0)?;
if alias.end != alias_text.len() {
return None;
}
Some(Some(alias.value))
}
fn strip_leading_as(text: &str) -> Option<&str> {
let prefix = text.get(..2)?;
if !prefix.eq_ignore_ascii_case("AS") {
return None;
}
let rest = &text[2..];
if rest.chars().next().is_some_and(char::is_whitespace) {
Some(rest)
} else {
None
}
}
fn is_select_star(body: &str, alias: Option<&str>) -> bool {
let trimmed = body.trim();
if trimmed == "*" {
return true;
}
let Some(alias) = alias else {
return false;
};
let Some((prefix, suffix)) = trimmed.split_once('.') else {
return false;
};
prefix.trim().eq_ignore_ascii_case(alias) && suffix.trim() == "*"
}
fn parse_from_source(body: &str) -> Option<FromSource> {
if body.is_empty()
|| body.contains(',')
|| body.contains('(')
|| body.contains(')')
|| contains_keyword(body, "JOIN")
{
return None;
}
let ident = parse_qualified_identifier(body)?;
if ident.parts.is_empty() || ident.parts.len() > 2 {
return None;
}
let tail = body[ident.end..].trim();
let alias = if tail.is_empty() {
None
} else {
let alias_text = strip_leading_as(tail).unwrap_or(tail).trim();
let alias_ident = read_identifier(alias_text, 0)?;
if alias_ident.end != alias_text.len() {
return None;
}
Some(alias_ident.value)
};
let table_name = ident.parts.last()?.clone();
let schema = if ident.parts.len() == 2 { Some(ident.parts[0].clone()) } else { None };
Some(FromSource { schema, table_name, alias })
}
fn parse_qualified_identifier(text: &str) -> Option<QualifiedIdentifier> {
let mut parts = Vec::new();
let mut pos = 0usize;
while pos < text.len() {
pos = skip_whitespace(text, pos);
let Some(ident) = read_identifier(text, pos) else {
break;
};
parts.push(ident.value);
pos = skip_whitespace(text, ident.end);
if !text[pos..].starts_with('.') {
break;
}
pos += 1;
}
if parts.is_empty() {
return None;
}
Some(QualifiedIdentifier { parts, end: pos })
}
fn read_identifier(text: &str, start: usize) -> Option<Identifier> {
let pos = skip_whitespace(text, start);
let mut chars = text[pos..].char_indices();
let (_, first) = chars.next()?;
if matches!(first, '"' | '`' | '[') {
let close = if first == '[' { ']' } else { first };
let mut value = String::new();
for (offset, ch) in chars {
if ch == close {
return Some(Identifier { value, end: pos + offset + ch.len_utf8() });
}
value.push(ch);
}
return None;
}
if !(first.is_ascii_alphabetic() || first == '_') {
return None;
}
let mut end = pos + first.len_utf8();
for (offset, ch) in text[end..].char_indices() {
if !(ch.is_ascii_alphanumeric() || ch == '_' || ch == '$') {
return Some(Identifier { value: text[pos..end + offset].to_string(), end: end + offset });
}
}
end = text.len();
Some(Identifier { value: text[pos..end].to_string(), end })
}
fn skip_whitespace(text: &str, pos: usize) -> usize {
let mut current = pos;
for (offset, ch) in text[pos..].char_indices() {
if !ch.is_whitespace() {
return pos + offset;
}
current = pos + offset + ch.len_utf8();
}
current
}
fn strip_sql_comments(sql: &str) -> String {
let mut result = String::new();
let mut chars = sql.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '-' && chars.peek() == Some(&'-') {
chars.next();
for next in chars.by_ref() {
if next == '\n' {
result.push('\n');
break;
}
}
continue;
}
if ch == '/' && chars.peek() == Some(&'*') {
chars.next();
let mut previous = '\0';
for next in chars.by_ref() {
if previous == '*' && next == '/' {
break;
}
previous = next;
}
continue;
}
result.push(ch);
}
result
}
fn has_top_level_keyword(sql: &str, keywords: &[&str]) -> bool {
keywords.iter().any(|keyword| find_top_level_keyword(sql, keyword, 0).is_some())
}
fn first_top_level_keyword_index(sql: &str, keywords: &[&str], start: usize) -> Option<usize> {
keywords.iter().filter_map(|keyword| find_top_level_keyword(sql, keyword, start)).min()
}
fn find_top_level_keyword(sql: &str, keyword: &str, start: usize) -> Option<usize> {
let mut depth = 0i32;
let mut quote: Option<char> = None;
let upper_keyword = keyword.to_ascii_uppercase();
for (index, ch) in sql.char_indices().filter(|(index, _)| *index >= start) {
if let Some(close) = quote {
if ch == close {
quote = None;
}
continue;
}
match ch {
'\'' | '"' | '`' => {
quote = Some(ch);
continue;
}
'[' => {
quote = Some(']');
continue;
}
'(' => {
depth += 1;
continue;
}
')' => {
depth = 0.max(depth - 1);
continue;
}
_ => {}
}
if depth != 0 {
continue;
}
let Some(candidate) = sql.get(index..index + keyword.len()) else {
continue;
};
if candidate.to_ascii_uppercase() != upper_keyword {
continue;
}
let before = previous_char(sql, index);
let after = sql[index + keyword.len()..].chars().next();
if !before.is_some_and(is_identifier_char) && !after.is_some_and(is_identifier_char) {
return Some(index);
}
}
None
}
fn starts_with_keyword(sql: &str, keyword: &str) -> bool {
let trimmed = sql.trim_start();
let Some(candidate) = trimmed.get(..keyword.len()) else {
return false;
};
if !candidate.eq_ignore_ascii_case(keyword) {
return false;
}
!trimmed[keyword.len()..].chars().next().is_some_and(is_identifier_char)
}
fn contains_keyword(sql: &str, keyword: &str) -> bool {
find_top_level_keyword(sql, keyword, 0).is_some()
}
fn previous_char(text: &str, index: usize) -> Option<char> {
text[..index].chars().next_back()
}
fn is_identifier_char(ch: char) -> bool {
ch.is_ascii_alphanumeric() || ch == '_' || ch == '$'
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn recognizes_simple_single_table_select_as_editable() {
let result =
analyze_editable_query_editability("select id, name from public.users where active = true order by id");
assert_eq!(
result,
QueryEditability {
editable: true,
analysis: Some(EditableQueryInfo {
schema: Some("public".to_string()),
table_name: "users".to_string(),
table_alias: None,
select_star: false,
columns: vec![
EditableQueryColumn {
source_name: Some("id".to_string()),
result_name: "id".to_string(),
expression: "id".to_string(),
},
EditableQueryColumn {
source_name: Some("name".to_string()),
result_name: "name".to_string(),
expression: "name".to_string(),
},
],
}),
reason: None,
}
);
}
#[test]
fn recognizes_quoted_table_names_and_aliases() {
let result =
analyze_editable_query_editability(r#"SELECT u."id", u."full name" FROM "app schema"."user table" AS u"#);
assert_eq!(
result.analysis.unwrap(),
EditableQueryInfo {
schema: Some("app schema".to_string()),
table_name: "user table".to_string(),
table_alias: Some("u".to_string()),
select_star: false,
columns: vec![
EditableQueryColumn {
source_name: Some("id".to_string()),
result_name: "id".to_string(),
expression: r#"u."id""#.to_string(),
},
EditableQueryColumn {
source_name: Some("full name".to_string()),
result_name: "full name".to_string(),
expression: r#"u."full name""#.to_string(),
},
],
}
);
}
#[test]
fn keeps_select_star_empty_columns() {
assert_eq!(
analyze_editable_query("select * from users").unwrap(),
EditableQueryInfo {
schema: None,
table_name: "users".to_string(),
table_alias: None,
select_star: true,
columns: Vec::new(),
}
);
}
#[test]
fn reports_joined_query_as_complex_source() {
let result =
analyze_editable_query_editability("select u.id, o.total from users u join orders o on o.user_id = u.id");
assert_eq!(result.editable, false);
assert_eq!(result.reason, Some(QueryEditabilityReason::ComplexSource));
}
#[test]
fn reports_grouped_query_as_aggregation() {
let result = analyze_editable_query_editability("select id, count(*) as total from users group by id");
assert_eq!(result.editable, false);
assert_eq!(result.reason, Some(QueryEditabilityReason::Aggregation));
}
#[test]
fn keeps_single_table_expression_columns() {
let result = analyze_editable_query_editability(
"select iso3, year, country_name, ihli / gdp_pc as score from ihli_data",
);
assert_eq!(
result.analysis.unwrap().columns,
vec![
EditableQueryColumn {
source_name: Some("iso3".to_string()),
result_name: "iso3".to_string(),
expression: "iso3".to_string(),
},
EditableQueryColumn {
source_name: Some("year".to_string()),
result_name: "year".to_string(),
expression: "year".to_string(),
},
EditableQueryColumn {
source_name: Some("country_name".to_string()),
result_name: "country_name".to_string(),
expression: "country_name".to_string(),
},
EditableQueryColumn {
source_name: None,
result_name: "score".to_string(),
expression: "ihli / gdp_pc".to_string(),
},
]
);
}
#[test]
fn serializes_reason_values_like_frontend_union() {
let json = serde_json::to_value(not_editable(QueryEditabilityReason::SetOperation)).unwrap();
assert_eq!(json, serde_json::json!({ "editable": false, "reason": "set-operation" }));
}
}

View File

@ -124,6 +124,7 @@ async fn main() {
.route("/query/execute-script", post(routes::query::execute_script))
.route("/query/execute-in-transaction", post(routes::query::execute_in_transaction))
.route("/query/analyze-sql-references", post(routes::query::analyze_sql_references))
.route("/query/analyze-editability", post(routes::query::analyze_editable_query_editability))
.route("/query/cancel", post(routes::query::cancel_query))
.route("/query/close-session", post(routes::query::close_query_session))
// Redis

View File

@ -51,6 +51,12 @@ pub struct AnalyzeSqlReferencesRequest {
pub dialect: Option<String>,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AnalyzeEditableQueryRequest {
pub sql: String,
}
pub async fn execute_query(
State(state): State<Arc<WebState>>,
Json(req): Json<ExecuteQueryRequest>,
@ -187,3 +193,9 @@ pub async fn analyze_sql_references(
) -> Result<Json<dbx_core::sql_analysis::SqlReferenceAnalysis>, AppError> {
dbx_core::sql_analysis::analyze_sql_references(&req.sql, req.dialect.as_deref()).map(Json).map_err(AppError)
}
pub async fn analyze_editable_query_editability(
Json(req): Json<AnalyzeEditableQueryRequest>,
) -> Json<dbx_core::sql_editability::QueryEditability> {
Json(dbx_core::sql_editability::analyze_editable_query_editability(&req.sql))
}

View File

@ -75,6 +75,12 @@ test("evicting cached tab results releases multi-result payloads and sessions",
headers: { "Content-Type": "application/json" },
});
}
if (url === "/api/query/analyze-editability") {
return new Response(JSON.stringify({ editable: false, reason: "complex-source" }), {
status: 200,
headers: { "Content-Type": "application/json" },
});
}
return new Response("unexpected request", { status: 500 });
}) as typeof fetch;

View File

@ -0,0 +1,30 @@
import { readFileSync } from "node:fs";
import { strict as assert } from "node:assert";
import test from "node:test";
const apiSource = readFileSync("apps/desktop/src/lib/api.ts", "utf8");
const tauriSource = readFileSync("apps/desktop/src/lib/tauri.ts", "utf8");
const httpSource = readFileSync("apps/desktop/src/lib/http.ts", "utf8");
const queryStoreSource = readFileSync("apps/desktop/src/stores/queryStore.ts", "utf8");
const tauriLibSource = readFileSync("src-tauri/src/lib.rs", "utf8");
const webMainSource = readFileSync("crates/dbx-web/src/main.rs", "utf8");
const rustCoreLibSource = readFileSync("crates/dbx-core/src/lib.rs", "utf8");
test("shared API exposes backend SQL editability analysis", () => {
assert.match(apiSource, /export const analyzeEditableQueryEditability = forward\("analyzeEditableQueryEditability"\)/);
assert.match(tauriSource, /export async function analyzeEditableQueryEditability\(/);
assert.match(tauriSource, /invoke\("analyze_editable_query_editability"/);
assert.match(httpSource, /export async function analyzeEditableQueryEditability\(/);
assert.match(httpSource, /\/api\/query\/analyze-editability/);
});
test("query metadata analysis uses backend SQL editability analysis", () => {
assert.match(queryStoreSource, /await api\.analyzeEditableQueryEditability\(sql\)/);
assert.doesNotMatch(queryStoreSource, /analyzeEditableQueryEditability,\n\s+sourceColumnsForResult/);
});
test("Rust backends register SQL editability analysis", () => {
assert.match(rustCoreLibSource, /pub mod sql_editability/);
assert.match(tauriLibSource, /commands::query::analyze_editable_query_editability/);
assert.match(webMainSource, /\/query\/analyze-editability/);
});

View File

@ -154,3 +154,8 @@ pub async fn analyze_sql_references(
) -> Result<dbx_core::sql_analysis::SqlReferenceAnalysis, String> {
dbx_core::sql_analysis::analyze_sql_references(&sql, dialect.as_deref())
}
#[tauri::command]
pub fn analyze_editable_query_editability(sql: String) -> Result<dbx_core::sql_editability::QueryEditability, String> {
Ok(dbx_core::sql_editability::analyze_editable_query_editability(&sql))
}

View File

@ -288,6 +288,7 @@ pub fn run() {
commands::query::execute_script,
commands::query::execute_in_transaction,
commands::query::analyze_sql_references,
commands::query::analyze_editable_query_editability,
commands::sql_file::preview_sql_file,
commands::sql_file::execute_sql_file,
commands::sql_file::cancel_sql_file_execution,