fix: support mysql install scripts without selected database (#2356)

This commit is contained in:
onenewcode
2026-07-02 13:18:49 +08:00
committed by GitHub
parent 730f839e74
commit 71e618dd46
12 changed files with 975 additions and 90 deletions
@@ -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";
@@ -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", () => {
@@ -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;
@@ -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<ConnectionConfig, "db_type" | "driver_profile">;
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;
}
+1 -1
View File
@@ -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<string> {
+1
View File
@@ -1641,6 +1641,7 @@ export interface SqlFilePreview {
filePath: string;
sizeBytes: number;
preview: string;
canExecuteWithoutSelectedDatabase: boolean;
}
export interface SqlFileProgress {
+29 -14
View File
@@ -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> {
+537 -59
View File
@@ -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<bool, String>,
}
struct MySqlSqlFileExecutor {
connection_id: String,
database: String,
pool_key: String,
db_type: Option<DatabaseType>,
bare: bool,
dialect: db::mysql::MySqlQueryDialect,
budget: DbOperationBudget,
conn: Option<mysql_async::Conn>,
}
impl MySqlSqlFileExecutor {
async fn build(
state: &AppState,
request: &SqlFileRequest,
import_target: Option<&SqlFileImportTarget>,
) -> Result<Option<Self>, 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<QueryResult, String> {
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<QueryResult, String> {
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<String> {
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(&quote) {
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<bool, String> {
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<bool, String> {
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<QueryResult, String> {
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);
}
}
+110
View File
@@ -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::<u16>().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")]]);
}
+1
View File
@@ -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),
})));
}
@@ -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<ConnectionConfig, "db_type" | "driver_profile"> {
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);
});
+4 -1
View File
@@ -32,13 +32,16 @@ pub async fn preview_sql_file(file_path: String) -> Result<SqlFilePreview, Strin
let path = PathBuf::from(&file_path);
let metadata = tokio::fs::metadata(&path).await.map_err(|e| e.to_string())?;
let bytes = tokio::fs::read(&path).await.map_err(|e| e.to_string())?;
let preview = decode_sql_file_bytes(&bytes)?.chars().take(20_000).collect();
let content = decode_sql_file_bytes(&bytes)?;
let preview = content.chars().take(20_000).collect();
Ok(SqlFilePreview {
file_name: path.file_name().and_then(|name| name.to_str()).unwrap_or("script.sql").to_string(),
file_path,
size_bytes: metadata.len(),
preview,
can_execute_without_selected_database:
dbx_core::sql_file_import::mysql_like_sql_file_can_execute_without_selected_database(&content),
})
}