From 71e618dd466ebad7f3a00102f3cbb908825b01a6 Mon Sep 17 00:00:00 2001 From: onenewcode Date: Thu, 2 Jul 2026 13:18:49 +0800 Subject: [PATCH] fix: support mysql install scripts without selected database (#2356) --- .../sql-file/SqlFileExecutionDialog.vue | 7 +- .../__tests__/useSqlExecution.spec.ts | 16 +- .../src/composables/useSqlExecution.ts | 13 +- .../lib/connectionLevelDatabaseBootstrap.ts | 217 +++++++ apps/desktop/src/lib/createDatabaseSql.ts | 2 +- apps/desktop/src/lib/tauri.ts | 1 + crates/dbx-core/src/sql.rs | 43 +- crates/dbx-core/src/sql_file_import.rs | 596 ++++++++++++++++-- crates/dbx-core/tests/live_mysql57.rs | 110 ++++ crates/dbx-web/src/routes/sql_file.rs | 1 + .../connectionLevelDatabaseBootstrap.test.ts | 54 ++ src-tauri/src/commands/sql_file.rs | 5 +- 12 files changed, 975 insertions(+), 90 deletions(-) create mode 100644 apps/desktop/src/lib/connectionLevelDatabaseBootstrap.ts create mode 100644 packages/app-tests/connectionLevelDatabaseBootstrap.test.ts diff --git a/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue b/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue index aff434a771..ca1f49019d 100644 --- a/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue +++ b/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue @@ -13,6 +13,7 @@ import DatabaseIcon from "@/components/icons/DatabaseIcon.vue"; import { useToast } from "@/composables/useToast"; import { useConnectionStore } from "@/stores/connectionStore"; import { databaseOptionsForConnection } from "@/composables/useDatabaseOptions"; +import { requiresSqlFileTargetDatabaseSelection } from "@/lib/connectionLevelDatabaseBootstrap"; import { cancelSqlFileExecution, executeSqlFile, listenSqlFileProgress, listDatabases, previewSqlFile, type SqlFilePreview, type SqlFileProgress, type SqlFileStatus } from "@/lib/api"; import { useExportTracker } from "@/composables/useExportTracker"; import { Check, CheckSquare, FileCode, FolderOpen, Loader2, Play, Square, X } from "@lucide/vue"; @@ -55,7 +56,11 @@ const sqlConnections = computed(() => store.connections.filter((c) => !["redis", const selectedConnection = computed(() => sqlConnections.value.find((c) => c.id === connectionId.value)); -const canStart = computed(() => Boolean(preview.value && selectedConnection.value && database.value.trim() && !running.value && !loadingPreview.value && !loadingDatabases.value)); +const canStart = computed(() => { + const connection = selectedConnection.value; + if (!preview.value || !connection || running.value || loadingPreview.value || loadingDatabases.value) return false; + return !!database.value.trim() || !requiresSqlFileTargetDatabaseSelection(connection, preview.value.canExecuteWithoutSelectedDatabase); +}); const statusTone = computed(() => { if (terminalStatus.value === "done") return "text-green-600"; diff --git a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts index 2c845ced34..a68f5abde9 100644 --- a/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSqlExecution.spec.ts @@ -39,8 +39,20 @@ describe("requiresDatabaseSelection", () => { expect(requiresDatabaseSelection(queryTab(), connection("mysql"), "CREATE SCHEMA `app-db` DEFAULT CHARACTER SET utf8mb4")).toBe(false); }); - it("requires a database when a MySQL batch includes non-creation statements", () => { - expect(requiresDatabaseSelection(queryTab(), connection("mysql"), "CREATE DATABASE app_db; SELECT * FROM users")).toBe(true); + it("allows MySQL install batches that switch databases before table DDL", () => { + expect(requiresDatabaseSelection(queryTab(), connection("mysql"), "CREATE DATABASE app_db; USE app_db; CREATE TABLE users(id INT PRIMARY KEY)")).toBe(false); + }); + + it("allows MySQL install batches with session setup before switching databases", () => { + expect(requiresDatabaseSelection(queryTab(), connection("mysql"), "SET NAMES utf8mb4; DROP DATABASE IF EXISTS app_db; CREATE DATABASE app_db; USE app_db; INSERT INTO users VALUES (1)")).toBe(false); + }); + + it("requires a database when MySQL batch statements never establish database context", () => { + expect(requiresDatabaseSelection(queryTab(), connection("mysql"), "CREATE DATABASE app_db; CREATE TABLE users(id INT)")).toBe(true); + }); + + it("requires a database when a USE statement is not a standalone database switch", () => { + expect(requiresDatabaseSelection(queryTab(), connection("mysql"), "CREATE DATABASE app_db; USE app_db SELECT 1; CREATE TABLE users(id INT)")).toBe(true); }); it("still requires a database for ordinary MySQL queries", () => { diff --git a/apps/desktop/src/composables/useSqlExecution.ts b/apps/desktop/src/composables/useSqlExecution.ts index 664241800a..1535e55cb6 100644 --- a/apps/desktop/src/composables/useSqlExecution.ts +++ b/apps/desktop/src/composables/useSqlExecution.ts @@ -6,6 +6,7 @@ import { useConnectionStore } from "@/stores/connectionStore"; import { useSettingsStore } from "@/stores/settingsStore"; import { useToast } from "@/composables/useToast"; import { isSingleDatabase } from "@/lib/databaseCapabilities"; +import { canExecuteWithoutSelectedDatabase } from "@/lib/connectionLevelDatabaseBootstrap"; import { classifySqlActivityKind } from "@/lib/historyActivityKind"; import { sqlMetadataRefreshTarget } from "@/lib/sqlMetadataRefresh"; import { classifyRedisCommandSafety, firstRedisCommandToken } from "@/lib/redisCommandSafety"; @@ -14,7 +15,6 @@ import { extractSqlParameters } from "@/lib/sqlParameters"; import type { ConnectionConfig, QueryTab } from "@/types/database"; const DANGER_RE = /^\s*(DROP|DELETE|TRUNCATE|ALTER|UPDATE|MERGE|REPLACE)\b/i; -const CONNECTION_LEVEL_CREATE_DATABASE_RE = /^\s*CREATE\s+(?:DATABASE|SCHEMA)\s+(?:IF\s+NOT\s+EXISTS\s+)?[`"[\]\w$.-]+/i; export function stripSqlComments(sql: string): string { return sql @@ -227,17 +227,6 @@ function supportsSqlTemplateParameters(connection: ConnectionConfig | undefined) return connection.db_type !== "redis" && connection.db_type !== "mongodb"; } -function canExecuteWithoutSelectedDatabase(connection: ConnectionConfig, sql: string): boolean { - if (connection.db_type !== "mysql") return false; - const statements = stripSqlComments(sql) - .split(";") - .map((stmt) => stmt.trim()) - .filter(Boolean); - // Only connection-level database creation may bypass the selected-database guard. - if (statements.length !== 1) return false; - return CONNECTION_LEVEL_CREATE_DATABASE_RE.test(statements[0]); -} - export function requiresDatabaseSelection(tab: QueryTab, connection: ConnectionConfig | undefined, sql = ""): boolean { if (tab.mode !== "query") return false; if (!connection || tab.database) return false; diff --git a/apps/desktop/src/lib/connectionLevelDatabaseBootstrap.ts b/apps/desktop/src/lib/connectionLevelDatabaseBootstrap.ts new file mode 100644 index 0000000000..3fb8b8ebf0 --- /dev/null +++ b/apps/desktop/src/lib/connectionLevelDatabaseBootstrap.ts @@ -0,0 +1,217 @@ +import { splitSqlStatementRanges } from "@/lib/sqlStatementRanges"; +import { supportsCreateDatabaseCharset } from "@/lib/createDatabaseSql"; +import type { ConnectionConfig, DatabaseType } from "@/types/database"; + +const MYSQL_BOOTSTRAP_EXTRA_PROFILES = new Set(["selectdb", "goldendb"]); + +type BootstrapConnection = Pick; +type ParsedIdentifier = { identifier: string; rest: string }; + +function splitBootstrapStatements(sql: string, dbType?: DatabaseType): string[] { + return splitSqlStatementRanges(sql, dbType) + .map((statement) => statement.sql.trim()) + .filter(Boolean); +} + +function parseMysqlIdentifier(input: string): ParsedIdentifier | null { + const first = input[0]; + if (!first) return null; + + if (first === "`" || first === '"') { + return parseDoubledDelimitedIdentifier(input, first); + } + if (first === "[") { + return parseBracketIdentifier(input); + } + + const end = input.search(/[\s;]/); + const identifier = (end === -1 ? input : input.slice(0, end)).trim(); + if (!identifier) return null; + return { identifier, rest: end === -1 ? "" : input.slice(end) }; +} + +function parseDoubledDelimitedIdentifier(input: string, quote: "`" | '"'): ParsedIdentifier | null { + let index = 1; + let segmentStart = 1; + let identifier = ""; + + while (index < input.length) { + if (input[index] === quote) { + identifier += input.slice(segmentStart, index); + if (input[index + 1] === quote) { + identifier += quote; + index += 2; + segmentStart = index; + continue; + } + return { identifier, rest: input.slice(index + 1) }; + } + index += 1; + } + + return null; +} + +function parseBracketIdentifier(input: string): ParsedIdentifier | null { + let index = 1; + let segmentStart = 1; + let identifier = ""; + + while (index < input.length) { + if (input[index] === "]") { + identifier += input.slice(segmentStart, index); + if (input[index + 1] === "]") { + identifier += "]"; + index += 2; + segmentStart = index; + continue; + } + return { identifier, rest: input.slice(index + 1) }; + } + index += 1; + } + + return null; +} + +function stripLeadingSqlComments(sql: string): string { + let rest = sql; + + for (;;) { + rest = rest.trimStart(); + if (!rest) return rest; + + if (rest.startsWith("--")) { + const index = rest.indexOf("\n"); + if (index === -1) return ""; + rest = rest.slice(index + 1); + continue; + } + + if (rest.startsWith("#")) { + const index = rest.indexOf("\n"); + if (index === -1) return ""; + rest = rest.slice(index + 1); + continue; + } + + if (rest.startsWith("/*")) { + const index = rest.indexOf("*/"); + if (index === -1) return ""; + rest = rest.slice(index + 2); + continue; + } + + return rest; + } +} + +function leadingSqlKeyword(statement: string): [keyword: string, rest: string] | null { + const trimmed = statement.trimStart(); + const match = trimmed.match(/^[A-Za-z_][A-Za-z0-9_]*/); + if (!match) return null; + return [match[0], trimmed.slice(match[0].length)]; +} + +function sqlRemainderIsCommentOnly(remainder: string): boolean { + let rest = remainder; + + for (;;) { + rest = rest.trimStart(); + if (!rest) return true; + + if (rest.startsWith(";")) { + rest = rest.slice(1); + continue; + } + + if (rest.startsWith("--") || rest.startsWith("#")) { + return true; + } + + if (rest.startsWith("/*")) { + const index = rest.indexOf("*/"); + if (index === -1) return false; + rest = rest.slice(index + 2); + continue; + } + + return false; + } +} + +function parseMysqlUseDatabaseTarget(statement: string): string | null { + const stripped = stripLeadingSqlComments(statement); + const match = stripped.match(/^USE\b/i); + if (!match) return null; + + const parsed = parseMysqlIdentifier(stripped.slice(match[0].length).trimStart()); + if (!parsed) return null; + + return sqlRemainderIsCommentOnly(parsed.rest) ? parsed.identifier : null; +} + +function isDatabaseKeyword(keyword: string): boolean { + const normalized = keyword.toUpperCase(); + return normalized === "DATABASE" || normalized === "SCHEMA"; +} + +function isAllowedBootstrapPrelude(statement: string): boolean { + const stripped = stripLeadingSqlComments(statement); + const leading = leadingSqlKeyword(stripped); + if (!leading) return false; + + const [keyword, rest] = leading; + const normalized = keyword.toUpperCase(); + if (normalized === "SET") { + return true; + } + + if (normalized !== "CREATE" && normalized !== "DROP") { + return false; + } + + const next = leadingSqlKeyword(rest); + return !!next && isDatabaseKeyword(next[0]); +} + +export function supportsConnectionLevelDatabaseBootstrap(connection: BootstrapConnection | undefined): boolean { + if (!connection) return false; + if (supportsCreateDatabaseCharset(connection.db_type, connection.driver_profile)) return true; + return !!connection.driver_profile && MYSQL_BOOTSTRAP_EXTRA_PROFILES.has(connection.driver_profile.toLowerCase()); +} + +export function canExecuteWithoutSelectedDatabase(connection: BootstrapConnection | undefined, sql: string): boolean { + if (!supportsConnectionLevelDatabaseBootstrap(connection)) return false; + + const statements = splitBootstrapStatements(sql, connection?.db_type); + if (statements.length === 0) return false; + + let sawStatement = false; + let hasDatabaseContext = false; + + for (const statement of statements) { + const stripped = stripLeadingSqlComments(statement); + if (!stripped) continue; + + sawStatement = true; + + if (isAllowedBootstrapPrelude(stripped)) { + continue; + } + if (parseMysqlUseDatabaseTarget(stripped)) { + hasDatabaseContext = true; + continue; + } + if (!hasDatabaseContext) { + return false; + } + } + + return sawStatement; +} + +export function requiresSqlFileTargetDatabaseSelection(connection: BootstrapConnection | undefined, sqlCanExecuteWithoutSelectedDatabase = false): boolean { + if (!supportsConnectionLevelDatabaseBootstrap(connection)) return true; + return !sqlCanExecuteWithoutSelectedDatabase; +} diff --git a/apps/desktop/src/lib/createDatabaseSql.ts b/apps/desktop/src/lib/createDatabaseSql.ts index c80e931993..f4708fe357 100644 --- a/apps/desktop/src/lib/createDatabaseSql.ts +++ b/apps/desktop/src/lib/createDatabaseSql.ts @@ -13,7 +13,7 @@ export interface CreateDatabaseSqlOptions { } export function supportsCreateDatabaseCharset(databaseType?: DatabaseType, driverProfile?: string | null): boolean { - return MYSQL_COMPATIBLE_TYPES.has(databaseType as DatabaseType) || (!!driverProfile && MYSQL_COMPATIBLE_PROFILES.has(driverProfile)); + return MYSQL_COMPATIBLE_TYPES.has(databaseType as DatabaseType) || (!!driverProfile && MYSQL_COMPATIBLE_PROFILES.has(driverProfile.toLowerCase())); } export function buildCreateDatabaseSql(options: CreateDatabaseSqlOptions): Promise { diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index bfd2e4dac6..f0ecb59d49 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -1641,6 +1641,7 @@ export interface SqlFilePreview { filePath: string; sizeBytes: number; preview: string; + canExecuteWithoutSelectedDatabase: boolean; } export interface SqlFileProgress { diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index f5626a24f8..ece8991b77 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -68,6 +68,7 @@ pub struct SqlFilePreview { pub file_path: String, pub size_bytes: u64, pub preview: String, + pub can_execute_without_selected_database: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -1478,6 +1479,31 @@ fn executable_sql_keyword_matches(token: &str, keyword: &str) -> bool { || (keyword.eq_ignore_ascii_case("DESCRIBE") && token.eq_ignore_ascii_case("DESC")) } +fn is_mysql_compatible_import_profile(profile: &str) -> bool { + matches!( + profile, + "mariadb" + | "tidb" + | "oceanbase" + | "custom_mysql" + | "doris" + | "starrocks" + | "manticoresearch" + | "selectdb" + | "goldendb" + ) +} + +pub(crate) fn supports_connection_level_database_bootstrap_target( + db_type: &DatabaseType, + driver_profile: Option<&str>, +) -> bool { + matches!(db_type, DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::Goldendb) + || driver_profile + .map(|profile| profile.to_ascii_lowercase()) + .is_some_and(|profile| is_mysql_compatible_import_profile(&profile) && profile != "manticoresearch") +} + fn is_mysql_compatible_import_target(db_type: &DatabaseType, driver_profile: Option<&str>) -> bool { matches!( db_type, @@ -1486,20 +1512,9 @@ fn is_mysql_compatible_import_target(db_type: &DatabaseType, driver_profile: Opt | DatabaseType::StarRocks | DatabaseType::ManticoreSearch | DatabaseType::Goldendb - ) || driver_profile.map(|profile| profile.to_ascii_lowercase()).is_some_and(|profile| { - matches!( - profile.as_str(), - "mariadb" - | "tidb" - | "oceanbase" - | "custom_mysql" - | "doris" - | "starrocks" - | "manticoresearch" - | "selectdb" - | "goldendb" - ) - }) + ) || driver_profile + .map(|profile| profile.to_ascii_lowercase()) + .is_some_and(|profile| is_mysql_compatible_import_profile(&profile)) } fn mysql_executable_comment_body(statement: &str) -> Option<&str> { diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 1857697897..7e766cda76 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -2,12 +2,17 @@ use std::time::Instant; use tokio_util::sync::CancellationToken; -use crate::connection::AppState; +use crate::connection::{AppState, PoolKind}; +use crate::db; use crate::models::connection::DatabaseType; -use crate::query::{execute_sql_statement_with_options, QueryExecutionOptions}; +use crate::query::{ + execute_sql_statement_with_options, pool_error_action, wait_for_query_opt, DbOperationBudget, PoolErrorAction, + QueryExecutionOptions, +}; use crate::sql::{ - optimize_sql_file_import_statements, statement_summary, SqlFileImportStatement, SqlFileImportStatementKind, - SqlFileProgress, SqlFileRequest, SqlFileStatus, SqlParsingOptions, SqlStatementSplitter, + optimize_sql_file_import_statements, prepare_sql_file_statement, statement_summary, SqlFileImportStatement, + SqlFileImportStatementKind, SqlFileProgress, SqlFileRequest, SqlFileStatementAction, SqlFileStatus, + SqlParsingOptions, SqlStatementSplitter, }; use crate::types::QueryResult; @@ -24,6 +29,177 @@ struct StatementErrorDecision { result: Result, } +struct MySqlSqlFileExecutor { + connection_id: String, + database: String, + pool_key: String, + db_type: Option, + bare: bool, + dialect: db::mysql::MySqlQueryDialect, + budget: DbOperationBudget, + conn: Option, +} + +impl MySqlSqlFileExecutor { + async fn build( + state: &AppState, + request: &SqlFileRequest, + import_target: Option<&SqlFileImportTarget>, + ) -> Result, String> { + let Some(target) = import_target else { + return Ok(None); + }; + if !crate::sql::supports_connection_level_database_bootstrap_target( + &target.db_type, + target.driver_profile.as_deref(), + ) { + return Ok(None); + } + + let database = request.database.trim(); + let database = (!database.is_empty()).then_some(database); + let pool_key = state.get_or_create_pool_for_session(&request.connection_id, database, None).await?; + let (db_type, driver_profile, bare) = { + let connections = state.connections.read().await; + let Some(PoolKind::Mysql(_, mode)) = connections.get(&pool_key) else { + return Ok(None); + }; + (Some(target.db_type), target.driver_profile.as_deref(), *mode == crate::connection::MysqlMode::Bare) + }; + let budget = { + let configs = state.configs.read().await; + let config = configs.get(&request.connection_id).ok_or("Connection config not found")?; + DbOperationBudget::from_connection_config(config) + }; + + Ok(Some(Self { + connection_id: request.connection_id.clone(), + database: request.database.clone(), + pool_key, + db_type, + bare, + dialect: db::mysql::MySqlQueryDialect::for_connection( + db_type.unwrap_or(DatabaseType::Mysql), + driver_profile, + ), + budget, + conn: None, + })) + } + + async fn execute_statement( + &mut self, + state: &AppState, + request: &SqlFileRequest, + sql: &str, + token: &CancellationToken, + statement_index: usize, + ) -> Result { + let execution_id = sql_file_statement_execution_id(&request.execution_id, statement_index); + let registered = state.running_queries.register(execution_id.clone()); + let child_token = registered.token(); + let cancel_task = { + let parent_token = token.clone(); + let running_queries = state.running_queries.clone(); + let execution_id = execution_id.clone(); + tokio::spawn(async move { + parent_token.cancelled().await; + running_queries.cancel(&execution_id); + }) + }; + + let result = self.execute_statement_inner(state, sql, &child_token, &execution_id).await; + + cancel_task.abort(); + result + } + + async fn execute_statement_inner( + &mut self, + state: &AppState, + sql: &str, + child_token: &CancellationToken, + execution_id: &str, + ) -> Result { + self.ensure_conn(state, child_token).await?; + state.running_queries.set_pool_key(execution_id, self.pool_key.clone()); + state.touch_pool_activity(&self.pool_key).await; + let _activity_touch = state.pool_activity_touch(&self.pool_key); + + let conn = self.conn.as_mut().ok_or("MySQL SQL file executor is missing a connection".to_string())?; + let connection_id = conn.id(); + let kill_opts = conn.opts().clone(); + state.running_queries.register_interrupt(execution_id, move || { + let kill_opts = kill_opts.clone(); + tokio::spawn(async move { + if let Err(error) = db::mysql::kill_query_with_opts(kill_opts, connection_id).await { + log::warn!("Failed to cancel MySQL SQL file import query {connection_id}: {error}"); + } + }); + }); + + let result = wait_for_query_opt( + Some(child_token.clone()), + self.budget.query_timeout, + db::mysql::execute_query_on_conn_with_max_rows(conn, sql, self.bare, None, self.dialect), + ) + .await; + + if result.is_ok() { + // Reconnects should reopen the most recent `USE` target rather than + // the request's initial database value. + if let Some(database) = mysql_use_database_target(sql) { + self.database = database; + } + } + + if let Err(error) = result.as_ref() { + let action = pool_error_action(self.db_type, error); + if matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) { + self.conn.take(); + if matches!(action, PoolErrorAction::ReconnectAndRetry) { + let database = self.database.trim(); + let database = (!database.is_empty()).then_some(database); + self.pool_key = state.reconnect_pool_for_session(&self.connection_id, database, None).await?; + } else { + state.remove_pool_by_key(&self.pool_key).await; + } + } + } + + result + } + + async fn ensure_conn(&mut self, state: &AppState, token: &CancellationToken) -> Result<(), String> { + if self.conn.is_some() { + return Ok(()); + } + + let database = self.database.trim(); + let database = (!database.is_empty()).then_some(database); + self.pool_key = state.get_or_create_pool_for_session(&self.connection_id, database, None).await?; + let pool = { + let connections = state.connections.read().await; + match connections.get(&self.pool_key) { + Some(PoolKind::Mysql(pool, _)) => pool.clone(), + Some(_) => return Err("SQL file import expected a MySQL-compatible pooled connection".to_string()), + None => return Err("Connection not found".to_string()), + } + }; + + self.conn = Some( + db::mysql::get_conn_with_health_check_with_cancel( + &pool, + self.budget.checkout_timeout, + self.budget.cleanup_timeout, + Some(token), + ) + .await?, + ); + Ok(()) + } +} + pub async fn execute_sql_file_content( state: &AppState, request: &SqlFileRequest, @@ -32,11 +208,6 @@ pub async fn execute_sql_file_content( started_at: Instant, mut emit: impl FnMut(SqlFileProgress), ) -> Result<(), String> { - let mut statement_index = 0; - let mut success_count = 0; - let mut failure_count = 0; - let mut affected_rows = 0; - let import_target = sql_file_import_target(state, &request.connection_id).await; let options = import_target.as_ref().map(|target| SqlParsingOptions::for_database_type(target.db_type)).unwrap_or_default(); @@ -49,55 +220,19 @@ pub async fn execute_sql_file_content( import_target.as_ref().map(|target| target.db_type), import_target.as_ref().and_then(|target| target.driver_profile.as_deref()), ); - - for planned_statement in planned_statements { - if token.is_cancelled() { - emit(sql_file_progress( - &request.execution_id, - SqlFileStatus::Cancelled, - statement_index, - success_count, - failure_count, - affected_rows, - started_at, - "", - None, - )); - return Ok(()); - } - - let next_statement_index = statement_index + planned_statement.source_statement_count; - if execute_statement_with_progress( - state, - request, - &token, - started_at, - next_statement_index, - &planned_statement, - &mut success_count, - &mut failure_count, - &mut affected_rows, - &mut emit, - ) - .await? - { - return Ok(()); - } - statement_index = next_statement_index; - } - - emit(sql_file_progress( - &request.execution_id, - SqlFileStatus::Done, - statement_index, - success_count, - failure_count, - affected_rows, + // MySQL-family imports need one pinned connection so `USE` and session + // state survive across the whole file. + let mut mysql_executor = MySqlSqlFileExecutor::build(state, request, import_target.as_ref()).await?; + execute_planned_statements_with_progress( + state, + request, + &token, started_at, - "", - None, - )); - Ok(()) + &planned_statements, + mysql_executor.as_mut(), + &mut emit, + ) + .await } #[allow(clippy::too_many_arguments)] @@ -136,6 +271,259 @@ async fn sql_file_import_target(state: &AppState, connection_id: &str) -> Option .map(|config| SqlFileImportTarget { db_type: config.db_type, driver_profile: config.driver_profile.clone() }) } +pub fn mysql_like_sql_file_can_execute_without_selected_database(file_content: &str) -> bool { + let options = SqlParsingOptions::mysql_compatible(); + let mut splitter = SqlStatementSplitter::with_options(options); + let mut statements = splitter.push_chunk(file_content); + statements.extend(splitter.finish()); + + let mut saw_statement = false; + let mut has_database_context = false; + + for statement in statements { + let prepared = match prepare_sql_file_statement(&statement, &DatabaseType::Mysql, None) { + SqlFileStatementAction::Skip => continue, + SqlFileStatementAction::Execute(sql) => sql, + }; + let statement = strip_leading_sql_comments(&prepared, true).trim_start(); + if statement.is_empty() { + continue; + } + + // Keep preview gating aligned with the executor: setup statements may + // run before the script establishes its own database context. + saw_statement = true; + let Some((keyword, remainder)) = leading_sql_keyword(statement) else { + return false; + }; + + if keyword.eq_ignore_ascii_case("SET") { + continue; + } + + if mysql_use_database_target(statement).is_some() { + has_database_context = true; + continue; + } + + if (keyword.eq_ignore_ascii_case("DROP") || keyword.eq_ignore_ascii_case("CREATE")) + && leading_sql_keyword(remainder) + .is_some_and(|(next, _)| next.eq_ignore_ascii_case("DATABASE") || next.eq_ignore_ascii_case("SCHEMA")) + { + continue; + } + + if !has_database_context { + return false; + } + } + + saw_statement +} + +fn mysql_use_database_target(sql: &str) -> Option { + let sql = strip_leading_sql_comments(sql, true).trim_start(); + let rest = sql.get(..3).filter(|prefix| prefix.eq_ignore_ascii_case("USE")).and_then(|_| sql.get(3..))?; + if rest.is_empty() || !rest.as_bytes()[0].is_ascii_whitespace() { + return None; + } + + let (database, remainder) = parse_mysql_identifier(rest.trim_start())?; + sql_remainder_is_comment_only(remainder).then_some(database) +} + +fn strip_leading_sql_comments(mut sql: &str, supports_hash_line_comments: bool) -> &str { + loop { + sql = sql.trim_start(); + if sql.is_empty() { + return sql; + } + + if let Some(rest) = sql.strip_prefix("--") { + if let Some(idx) = rest.find('\n') { + sql = &rest[idx + 1..]; + continue; + } + return ""; + } + + if supports_hash_line_comments { + if let Some(rest) = sql.strip_prefix('#') { + if let Some(idx) = rest.find('\n') { + sql = &rest[idx + 1..]; + continue; + } + return ""; + } + } + + if let Some(rest) = sql.strip_prefix("/*") { + let Some(close) = rest.find("*/") else { + return ""; + }; + sql = &rest[close + 2..]; + continue; + } + + return sql; + } +} + +fn leading_sql_keyword(input: &str) -> Option<(&str, &str)> { + let input = input.trim_start(); + let end = input.find(|ch: char| !(ch.is_ascii_alphanumeric() || ch == '_')).unwrap_or(input.len()); + (end > 0).then_some((&input[..end], &input[end..])) +} + +fn parse_mysql_identifier(input: &str) -> Option<(String, &str)> { + let first = *input.as_bytes().first()?; + match first { + b'`' | b'"' => parse_mysql_doubled_delimited_identifier(input, first), + b'[' => parse_mysql_bracket_identifier(input), + _ => { + let end = input.find(|c: char| c.is_whitespace() || c == ';').unwrap_or(input.len()); + let identifier = input[..end].trim(); + (!identifier.is_empty()).then_some((identifier.to_string(), &input[end..])) + } + } +} + +fn parse_mysql_doubled_delimited_identifier(input: &str, quote: u8) -> Option<(String, &str)> { + let bytes = input.as_bytes(); + let mut index = 1; + let mut segment_start = 1; + let mut identifier = String::new(); + + while index < bytes.len() { + if bytes[index] == quote { + identifier.push_str(&input[segment_start..index]); + if bytes.get(index + 1) == Some("e) { + identifier.push(quote as char); + index += 2; + segment_start = index; + continue; + } + return Some((identifier, &input[index + 1..])); + } + index += 1; + } + + None +} + +fn parse_mysql_bracket_identifier(input: &str) -> Option<(String, &str)> { + let bytes = input.as_bytes(); + let mut index = 1; + let mut segment_start = 1; + let mut identifier = String::new(); + + while index < bytes.len() { + if bytes[index] == b']' { + identifier.push_str(&input[segment_start..index]); + if bytes.get(index + 1) == Some(&b']') { + identifier.push(']'); + index += 2; + segment_start = index; + continue; + } + return Some((identifier, &input[index + 1..])); + } + index += 1; + } + + None +} + +fn sql_remainder_is_comment_only(mut remainder: &str) -> bool { + loop { + remainder = remainder.trim_start(); + if remainder.is_empty() { + return true; + } + if let Some(rest) = remainder.strip_prefix(';') { + remainder = rest; + continue; + } + if remainder.starts_with("--") || remainder.starts_with('#') { + return true; + } + if let Some(rest) = remainder.strip_prefix("/*") { + let Some(close) = rest.find("*/") else { + return false; + }; + remainder = &rest[close + 2..]; + continue; + } + return false; + } +} + +#[allow(clippy::too_many_arguments)] +async fn execute_planned_statements_with_progress( + state: &AppState, + request: &SqlFileRequest, + token: &CancellationToken, + started_at: Instant, + planned_statements: &[SqlFileImportStatement], + mut mysql_executor: Option<&mut MySqlSqlFileExecutor>, + emit: &mut impl FnMut(SqlFileProgress), +) -> Result<(), String> { + let mut statement_index = 0; + let mut success_count = 0; + let mut failure_count = 0; + let mut affected_rows = 0; + + for planned_statement in planned_statements { + if token.is_cancelled() { + emit(sql_file_progress( + &request.execution_id, + SqlFileStatus::Cancelled, + statement_index, + success_count, + failure_count, + affected_rows, + started_at, + "", + None, + )); + return Ok(()); + } + + let next_statement_index = statement_index + planned_statement.source_statement_count; + if execute_statement_with_progress( + state, + request, + token, + started_at, + next_statement_index, + planned_statement, + &mut success_count, + &mut failure_count, + &mut affected_rows, + mysql_executor.as_deref_mut(), + emit, + ) + .await? + { + return Ok(()); + } + statement_index = next_statement_index; + } + + emit(sql_file_progress( + &request.execution_id, + SqlFileStatus::Done, + statement_index, + success_count, + failure_count, + affected_rows, + started_at, + "", + None, + )); + Ok(()) +} + #[allow(clippy::too_many_arguments)] async fn execute_statement_with_progress( state: &AppState, @@ -147,6 +535,7 @@ async fn execute_statement_with_progress( success_count: &mut usize, failure_count: &mut usize, affected_rows: &mut u64, + mut mysql_executor: Option<&mut MySqlSqlFileExecutor>, emit: &mut impl FnMut(SqlFileProgress), ) -> Result { if token.is_cancelled() { @@ -206,7 +595,13 @@ async fn execute_statement_with_progress( None, )); - match execute_sql_file_statement(state, request, &statement.sql, token, statement_index).await { + let result = { + let mysql_executor = mysql_executor.as_deref_mut(); + execute_sql_file_statement_with_executor(state, request, &statement.sql, token, statement_index, mysql_executor) + .await + }; + + match result { Ok(result) => { *success_count += statement.source_statement_count; *affected_rows += result.affected_rows; @@ -235,6 +630,7 @@ async fn execute_statement_with_progress( success_count, failure_count, affected_rows, + mysql_executor, emit, ) .await; @@ -273,6 +669,7 @@ async fn execute_merged_statement_fallback_with_progress( success_count: &mut usize, failure_count: &mut usize, affected_rows: &mut u64, + mut mysql_executor: Option<&mut MySqlSqlFileExecutor>, emit: &mut impl FnMut(SqlFileProgress), ) -> Result { for (offset, source_sql) in statement.source_sqls.iter().enumerate() { @@ -305,7 +702,16 @@ async fn execute_merged_statement_fallback_with_progress( None, )); - match execute_sql_file_statement(state, request, source_sql, token, statement_index).await { + match execute_sql_file_statement_with_executor( + state, + request, + source_sql, + token, + statement_index, + mysql_executor.as_deref_mut(), + ) + .await + { Ok(result) => { *success_count += 1; *affected_rows += result.affected_rows; @@ -349,6 +755,21 @@ async fn execute_merged_statement_fallback_with_progress( Ok(false) } +async fn execute_sql_file_statement_with_executor( + state: &AppState, + request: &SqlFileRequest, + sql: &str, + token: &CancellationToken, + statement_index: usize, + mysql_executor: Option<&mut MySqlSqlFileExecutor>, +) -> Result { + if let Some(mysql_executor) = mysql_executor { + mysql_executor.execute_statement(state, request, sql, token, statement_index).await + } else { + execute_sql_file_statement(state, request, sql, token, statement_index).await + } +} + async fn execute_sql_file_statement( state: &AppState, request: &SqlFileRequest, @@ -454,6 +875,7 @@ fn statement_error_decision( #[cfg(test)] mod tests { use super::*; + use crate::models::connection::DatabaseType; #[test] fn stop_on_error_returns_err_with_terminal_error_progress() { @@ -520,4 +942,60 @@ mod tests { assert_eq!(value["status"], "statementDone"); assert!(value.get("execution_id").is_none()); } + + #[test] + fn supports_connection_level_database_bootstrap_for_mysql_like_targets() { + assert!(crate::sql::supports_connection_level_database_bootstrap_target(&DatabaseType::Mysql, None)); + assert!(crate::sql::supports_connection_level_database_bootstrap_target(&DatabaseType::Doris, None)); + assert!(crate::sql::supports_connection_level_database_bootstrap_target(&DatabaseType::Goldendb, None)); + assert!(crate::sql::supports_connection_level_database_bootstrap_target( + &DatabaseType::Mysql, + Some("selectdb") + )); + assert!(crate::sql::supports_connection_level_database_bootstrap_target( + &DatabaseType::Mysql, + Some("oceanbase") + )); + } + + #[test] + fn excludes_non_mysql_bootstrap_targets() { + assert!(!crate::sql::supports_connection_level_database_bootstrap_target(&DatabaseType::Postgres, None)); + assert!( + !crate::sql::supports_connection_level_database_bootstrap_target(&DatabaseType::ManticoreSearch, None,) + ); + assert!( + !crate::sql::supports_connection_level_database_bootstrap_target(&DatabaseType::OceanbaseOracle, None,) + ); + } + + #[test] + fn mysql_like_sql_file_without_selected_database_requires_bootstrap_context() { + assert!(mysql_like_sql_file_can_execute_without_selected_database( + "SET NAMES utf8mb4;\nCREATE DATABASE app_db;\n-- switch tenant\nUSE app_db;\nCREATE TABLE users(id INT)" + )); + assert!(!mysql_like_sql_file_can_execute_without_selected_database( + "CREATE DATABASE app_db;\nCREATE TABLE users(id INT)" + )); + assert!(!mysql_like_sql_file_can_execute_without_selected_database( + "CREATE DATABASE app_db;\nUSE app_db SELECT 1;\nCREATE TABLE users(id INT)" + )); + } + + #[test] + fn parses_mysql_use_database_targets() { + assert_eq!(mysql_use_database_target("USE app_db"), Some("app_db".to_string())); + assert_eq!(mysql_use_database_target(" use `app-db` ; "), Some("app-db".to_string())); + assert_eq!(mysql_use_database_target(r#"USE "tenant""01""#), Some(r#"tenant"01"#.to_string())); + assert_eq!(mysql_use_database_target("USE [tenant]]01]"), Some("tenant]01".to_string())); + assert_eq!(mysql_use_database_target("USE app_db; -- switch tenant"), Some("app_db".to_string())); + assert_eq!(mysql_use_database_target("-- switch tenant\nUSE app_db"), Some("app_db".to_string())); + } + + #[test] + fn ignores_non_terminal_use_statements() { + assert_eq!(mysql_use_database_target("SELECT 1"), None); + assert_eq!(mysql_use_database_target("USE"), None); + assert_eq!(mysql_use_database_target("USE app_db SELECT 1"), None); + } } diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 1af2dab5cd..5e796f2f4e 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -1,3 +1,43 @@ +use dbx_core::connection::AppState; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::query::execute_sql_statement; +use dbx_core::sql::SqlFileRequest; +use dbx_core::sql_file_import::execute_sql_file_content; +use dbx_core::storage::Storage; +use tokio_util::sync::CancellationToken; + +fn live_mysql_sql_file_config(id: &str) -> ConnectionConfig { + let host = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_HOST").expect("DBX_LIVE_SQL_FILE_MYSQL_HOST"); + let port = + std::env::var("DBX_LIVE_SQL_FILE_MYSQL_PORT").ok().and_then(|value| value.parse::().ok()).unwrap_or(3306); + let username = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_USER").expect("DBX_LIVE_SQL_FILE_MYSQL_USER"); + let password = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_PASSWORD").expect("DBX_LIVE_SQL_FILE_MYSQL_PASSWORD"); + + serde_json::from_value(serde_json::json!({ + "id": id, + "name": id, + "db_type": DatabaseType::Mysql, + "host": host, + "port": port, + "username": username, + "password": password, + "database": null, + "connect_timeout_secs": 5, + "query_timeout_secs": 30, + "idle_timeout_secs": 60, + "keepalive_interval_secs": 0 + })) + .expect("live MySQL SQL file config should deserialize") +} + +async fn app_state_with_config(config: ConnectionConfig) -> (AppState, std::path::PathBuf) { + let db_path = std::env::temp_dir().join(format!("dbx-live-sql-file-{}.db", uuid::Uuid::new_v4().simple())); + let storage = Storage::open(&db_path).await.expect("open temp storage"); + let state = AppState::new(storage); + state.configs.write().await.insert(config.id.clone(), config); + (state, db_path) +} + #[tokio::test] #[ignore = "requires the remote DBX MySQL 5.7 smoke-test container"] async fn live_mysql57_text_protocol_select_succeeds() { @@ -190,3 +230,73 @@ async fn live_mysql_recovers_after_server_idle_disconnect() { assert_eq!(result.columns, vec!["recovered"]); assert_eq!(result.rows, vec![vec![serde_json::json!("1")]]); } + +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQL_FILE_MYSQL_* env vars pointing at a writable MySQL connection without a required default database"] +async fn live_sql_file_import_creates_database_and_switches_context_without_default_database() { + let config = live_mysql_sql_file_config("sql-file-import"); + let (state, db_path) = app_state_with_config(config.clone()).await; + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let database_name = format!("dbx_issue_2356_{suffix}"); + let script = format!( + r#" +CREATE DATABASE `{database_name}`; +USE `{database_name}`; +CREATE TABLE install_check ( + id INT PRIMARY KEY +); +INSERT INTO install_check (id) VALUES (1), (2); +"# + ); + let request = SqlFileRequest { + execution_id: format!("exec-{suffix}"), + connection_id: config.id.clone(), + database: String::new(), + file_path: "issue-2356-mysql-install.sql".to_string(), + continue_on_error: false, + }; + + let _ = execute_sql_statement( + &state, + &config.id, + "", + &format!("DROP DATABASE IF EXISTS `{database_name}`"), + None, + None, + ) + .await; + + let result = execute_sql_file_content( + &state, + &request, + &script, + CancellationToken::new(), + std::time::Instant::now(), + |_| {}, + ) + .await; + let verify = execute_sql_statement( + &state, + &config.id, + &database_name, + "SELECT COUNT(*) AS count FROM install_check", + None, + None, + ) + .await; + let _ = execute_sql_statement( + &state, + &config.id, + "", + &format!("DROP DATABASE IF EXISTS `{database_name}`"), + None, + None, + ) + .await; + let _ = std::fs::remove_file(db_path); + + result.expect("SQL file import should succeed"); + let verify = verify.expect("verify imported rows"); + assert_eq!(verify.columns, vec!["count"]); + assert_eq!(verify.rows, vec![vec![serde_json::json!("2")]]); +} diff --git a/crates/dbx-web/src/routes/sql_file.rs b/crates/dbx-web/src/routes/sql_file.rs index d1624645eb..4b416adbd5 100644 --- a/crates/dbx-web/src/routes/sql_file.rs +++ b/crates/dbx-web/src/routes/sql_file.rs @@ -52,6 +52,7 @@ pub async fn preview_sql_file( "filePath": file_path.to_string_lossy(), "sizeBytes": size_bytes, "preview": preview, + "canExecuteWithoutSelectedDatabase": dbx_core::sql_file_import::mysql_like_sql_file_can_execute_without_selected_database(&content), }))); } diff --git a/packages/app-tests/connectionLevelDatabaseBootstrap.test.ts b/packages/app-tests/connectionLevelDatabaseBootstrap.test.ts new file mode 100644 index 0000000000..2311a7231c --- /dev/null +++ b/packages/app-tests/connectionLevelDatabaseBootstrap.test.ts @@ -0,0 +1,54 @@ +import { strict as assert } from "node:assert"; +import { test } from "vitest"; +import { canExecuteWithoutSelectedDatabase, requiresSqlFileTargetDatabaseSelection, supportsConnectionLevelDatabaseBootstrap } from "../../apps/desktop/src/lib/connectionLevelDatabaseBootstrap.ts"; +import type { ConnectionConfig } from "../../apps/desktop/src/types/database.ts"; + +function connection(dbType: ConnectionConfig["db_type"], driverProfile?: string): Pick { + return { + db_type: dbType, + driver_profile: driverProfile, + }; +} + +test("supports MySQL-compatible connection-level bootstrap targets", () => { + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("mysql")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("doris")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("starrocks")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("goldendb")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("mysql", "selectdb")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("mysql", "oceanbase")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("mysql", "OceanBase")), true); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("mysql", "SelectDB")), true); +}); + +test("excludes non-MySQL bootstrap targets", () => { + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("postgres")), false); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("manticoresearch")), false); + assert.equal(supportsConnectionLevelDatabaseBootstrap(connection("oceanbase-oracle")), false); +}); + +test("allows install scripts that create and switch databases before table DDL", () => { + const sql = "SET NAMES utf8mb4; DROP DATABASE IF EXISTS app_db; CREATE DATABASE app_db; USE app_db; CREATE TABLE users(id INT)"; + assert.equal(canExecuteWithoutSelectedDatabase(connection("mysql"), sql), true); +}); + +test("blocks scripts that never establish database context", () => { + const sql = "CREATE DATABASE app_db; CREATE TABLE users(id INT)"; + assert.equal(canExecuteWithoutSelectedDatabase(connection("mysql"), sql), false); +}); + +test("blocks malformed USE statements that trail into other SQL", () => { + const sql = "CREATE DATABASE app_db; USE app_db SELECT 1; CREATE TABLE users(id INT)"; + assert.equal(canExecuteWithoutSelectedDatabase(connection("mysql"), sql), false); +}); + +test("accepts escaped MySQL-compatible identifiers in USE statements", () => { + const sql = "CREATE DATABASE [tenant]]01]; -- switch tenant\nUSE [tenant]]01]; CREATE TABLE users(id INT)"; + assert.equal(canExecuteWithoutSelectedDatabase(connection("mysql"), sql), true); +}); + +test("SQL file execution still requires an explicit database for non-MySQL targets", () => { + assert.equal(requiresSqlFileTargetDatabaseSelection(connection("postgres"), true), true); + assert.equal(requiresSqlFileTargetDatabaseSelection(connection("mysql"), true), false); + assert.equal(requiresSqlFileTargetDatabaseSelection(connection("mysql"), false), true); +}); diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index a12a769336..540d882f45 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -32,13 +32,16 @@ pub async fn preview_sql_file(file_path: String) -> Result