diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 39b3a8b72b..0a726360aa 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -4205,7 +4205,9 @@ const isDragging = computed(() => dragState.active && dragState.draggedId === pr const TABLE_REFERENCE_DRAG_THRESHOLD = 5; const TABLE_REFERENCE_DRAGGING_CLASS = "dbx-table-reference-dragging"; const canDragTableReference = computed(() => { - if (props.dragDisabled || !props.node.connectionId || props.node.database == null) return false; + if (props.dragDisabled || !props.node.connectionId) return false; + if (props.node.type === "database") return typeof props.node.database === "string" && props.node.database.trim().length > 0; + if (props.node.database == null) return false; if (props.node.type === "table" || props.node.type === "view" || props.node.type === "materialized_view") return true; return props.node.type === "column" && !!props.node.tableName; }); @@ -4220,6 +4222,14 @@ let suppressNextTableReferenceClick = false; function tableReferenceDragPayload(): QueryEditorTableReferencePayload | null { if (!canDragTableReference.value) return null; + if (props.node.type === "database") { + return createTableReferencePayload({ + connectionId: props.node.connectionId, + database: props.node.database, + referenceType: "database", + databaseType: currentDatabaseType(), + }); + } if (props.node.type === "column") { const columnName = columnNameForDrag(props.node); if (!props.node.tableName || !columnName) return null; diff --git a/apps/desktop/src/lib/__tests__/editor/queryEditorTableDrop.spec.ts b/apps/desktop/src/lib/__tests__/editor/queryEditorTableDrop.spec.ts new file mode 100644 index 0000000000..5a1209b787 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/editor/queryEditorTableDrop.spec.ts @@ -0,0 +1,27 @@ +import { describe, expect, it } from "vitest"; +import { createTableReferencePayload, parseTableReferencePayload, tableReferenceInsertText } from "@/lib/editor/queryEditorTableDrop"; + +describe("query editor table reference drop", () => { + it("inserts a quoted database name for database references", () => { + const payload = createTableReferencePayload({ + connectionId: "conn-1", + database: "app-db", + referenceType: "database", + databaseType: "mysql", + }); + + expect(payload).not.toBeNull(); + expect(tableReferenceInsertText(payload!)).toBe("`app-db`"); + }); + + it("round-trips database reference payloads", () => { + const payload = createTableReferencePayload({ + connectionId: "conn-1", + database: "reporting", + referenceType: "database", + databaseType: "postgres", + })!; + + expect(parseTableReferencePayload(JSON.stringify(payload))).toEqual(payload); + }); +}); diff --git a/apps/desktop/src/lib/editor/queryEditorTableDrop.ts b/apps/desktop/src/lib/editor/queryEditorTableDrop.ts index f1853c8ddb..71171ee098 100644 --- a/apps/desktop/src/lib/editor/queryEditorTableDrop.ts +++ b/apps/desktop/src/lib/editor/queryEditorTableDrop.ts @@ -9,9 +9,9 @@ export interface QueryEditorTableReferencePayload { connectionId: string; database: string; schema?: string; - tableName: string; + tableName?: string; columnName?: string; - referenceType?: "table" | "column"; + referenceType?: "database" | "table" | "column"; databaseType?: DatabaseType; } @@ -23,15 +23,21 @@ export interface QueryEditorTableReferenceDropDetail { let activeTableReferencePayload: QueryEditorTableReferencePayload | null = null; -export function createTableReferencePayload(options: { connectionId?: string; database?: string; schema?: string; tableName?: string; columnName?: string; databaseType?: DatabaseType }): QueryEditorTableReferencePayload | null { - if (!options.connectionId || options.database == null || !options.tableName) return null; +export function createTableReferencePayload(options: { connectionId?: string; database?: string; schema?: string; tableName?: string; columnName?: string; referenceType?: "database" | "table" | "column"; databaseType?: DatabaseType }): QueryEditorTableReferencePayload | null { + if (!options.connectionId || options.database == null) return null; + const referenceType = options.referenceType ?? (options.columnName ? "column" : "table"); + if (referenceType !== "database" && !options.tableName) return null; const payload: QueryEditorTableReferencePayload = { kind: "dbx-table-reference", connectionId: options.connectionId, database: options.database, - tableName: options.tableName, }; - if (options.columnName) { + if (referenceType === "database") { + payload.referenceType = "database"; + } else { + payload.tableName = options.tableName; + } + if (referenceType === "column" && options.columnName) { payload.columnName = options.columnName; payload.referenceType = "column"; } @@ -48,9 +54,20 @@ export function parseTableReferencePayload(value: string | undefined | null): Qu if (!value) return null; try { const parsed = JSON.parse(value) as Partial; - if (parsed.kind !== "dbx-table-reference" || typeof parsed.connectionId !== "string" || typeof parsed.database !== "string" || typeof parsed.tableName !== "string" || !parsed.connectionId || !parsed.tableName) { + if (parsed.kind !== "dbx-table-reference" || typeof parsed.connectionId !== "string" || typeof parsed.database !== "string" || !parsed.connectionId) { return null; } + if (parsed.referenceType === "database") { + const payload: QueryEditorTableReferencePayload = { + kind: "dbx-table-reference", + connectionId: parsed.connectionId, + database: parsed.database, + referenceType: "database", + }; + if (parsed.databaseType) payload.databaseType = parsed.databaseType; + return payload; + } + if (typeof parsed.tableName !== "string" || !parsed.tableName) return null; const columnName = typeof parsed.columnName === "string" && parsed.columnName ? parsed.columnName : undefined; const referenceType = parsed.referenceType === "column" || columnName ? "column" : "table"; if (referenceType === "column" && !columnName) return null; @@ -100,12 +117,16 @@ export function createTableReferenceDropEvent(detail: QueryEditorTableReferenceD export function tableReferenceInsertText(payload: QueryEditorTableReferencePayload, fallbackDatabaseType?: DatabaseType): string { const databaseType = payload.databaseType ?? fallbackDatabaseType; + if (payload.referenceType === "database") { + return quoteTableIdentifier(databaseType, payload.database); + } if (payload.referenceType === "column" && payload.columnName) { return quoteTableIdentifier(databaseType, payload.columnName); } + const tableName = payload.tableName || payload.database; return qualifiedTableName({ databaseType, schema: payload.schema, - tableName: payload.tableName, + tableName, }); } diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 69bb8c63ed..20f2803674 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -5,6 +5,7 @@ use std::io::Write; use tokio::sync::RwLock; use crate::models::connection::DatabaseType; +use crate::object_source_sql::build_export_object_source_sql; use crate::sql_dialect::{qualified_table_name, quote_table_identifier, uses_single_row_insert_statements}; use crate::transfer::{ format_ch_array_sql_literal, format_pg_array_sql_literal, is_identity_column_extra, quote_identifier, @@ -1224,8 +1225,10 @@ pub async fn export_database_sql_core( .await { Ok(obj_source) => { - if !obj_source.source.is_empty() { - writeln!(file, "{};\n", obj_source.source).map_err(|e| format!("Failed to write file: {e}"))?; + let source = + build_export_object_source_sql(db_type, crate::db::ObjectSourceKind::View, &obj_source.source); + if !source.is_empty() { + writeln!(file, "{source}\n").map_err(|e| format!("Failed to write file: {e}"))?; } } Err(e) => { @@ -1265,8 +1268,13 @@ pub async fn export_database_sql_core( .await { Ok(obj_source) => { - if !obj_source.source.is_empty() { - writeln!(file, "{};\n", obj_source.source).map_err(|e| format!("Failed to write file: {e}"))?; + let source = build_export_object_source_sql( + db_type, + crate::db::ObjectSourceKind::Procedure, + &obj_source.source, + ); + if !source.is_empty() { + writeln!(file, "{source}\n").map_err(|e| format!("Failed to write file: {e}"))?; } } Err(e) => { @@ -1306,8 +1314,13 @@ pub async fn export_database_sql_core( .await { Ok(obj_source) => { - if !obj_source.source.is_empty() { - writeln!(file, "{};\n", obj_source.source).map_err(|e| format!("Failed to write file: {e}"))?; + let source = build_export_object_source_sql( + db_type, + crate::db::ObjectSourceKind::Function, + &obj_source.source, + ); + if !source.is_empty() { + writeln!(file, "{source}\n").map_err(|e| format!("Failed to write file: {e}"))?; } } Err(e) => { diff --git a/crates/dbx-core/src/object_source_sql.rs b/crates/dbx-core/src/object_source_sql.rs index 778a5eff4c..f33b30c121 100644 --- a/crates/dbx-core/src/object_source_sql.rs +++ b/crates/dbx-core/src/object_source_sql.rs @@ -227,6 +227,21 @@ pub fn build_view_ddl_sql(input: BuildViewDdlInput) -> String { format!("CREATE VIEW {qualified_name} AS\n{}", ensure_semicolon(source)) } +pub fn build_export_object_source_sql( + database_type: DatabaseType, + object_type: ObjectSourceKind, + source: &str, +) -> String { + let source = source.trim(); + if source.is_empty() { + return String::new(); + } + if is_mysql_like(database_type) && matches!(object_type, ObjectSourceKind::Procedure | ObjectSourceKind::Function) { + return mysql_delimited_routine_source(source); + } + ensure_semicolon(source) +} + pub fn object_source_save_execution_mode(_database_type: DatabaseType) -> ObjectSourceSaveExecutionMode { ObjectSourceSaveExecutionMode::Single } @@ -317,6 +332,23 @@ fn ensure_semicolon(sql: &str) -> String { } } +fn mysql_delimited_routine_source(source: &str) -> String { + let trimmed = source.trim(); + if Regex::new(r"(?i)^\s*DELIMITER\b").unwrap().is_match(trimmed) { + return trimmed.to_string(); + } + let body = trimmed.trim_end_matches(';').trim_end(); + let delimiter = mysql_routine_script_delimiter(body); + format!("DELIMITER {delimiter}\n{body}{delimiter}\nDELIMITER ;") +} + +fn mysql_routine_script_delimiter(source: &str) -> &'static str { + ["//", "$$", ";;", "__DBX_DELIMITER__"] + .into_iter() + .find(|delimiter| !source.contains(delimiter)) + .unwrap_or("__DBX_DELIMITER__") +} + fn source_starts_with_create_or_alter(source: &str) -> bool { Regex::new(r"(?i)^\s*(?:CREATE|ALTER)\s+").unwrap().is_match(source) } @@ -1217,6 +1249,40 @@ mod tests { ); } + #[test] + fn mysql_routine_export_uses_delimiter_script() { + let sql = build_export_object_source_sql( + DatabaseType::Mysql, + ObjectSourceKind::Procedure, + "CREATE DEFINER=`root`@`%` PROCEDURE `refresh_cache`()\nBEGIN\n SELECT 1;\nEND", + ); + + assert_eq!( + sql, + "DELIMITER //\nCREATE DEFINER=`root`@`%` PROCEDURE `refresh_cache`()\nBEGIN\n SELECT 1;\nEND//\nDELIMITER ;" + ); + } + + #[test] + fn mysql_routine_export_does_not_double_wrap_delimiter_script() { + let source = "DELIMITER //\nCREATE PROCEDURE `refresh_cache`()\nBEGIN\n SELECT 1;\nEND//\nDELIMITER ;"; + + let sql = build_export_object_source_sql(DatabaseType::Mysql, ObjectSourceKind::Procedure, source); + + assert_eq!(sql, source); + } + + #[test] + fn mysql_view_export_keeps_regular_statement_terminator() { + let sql = build_export_object_source_sql( + DatabaseType::Mysql, + ObjectSourceKind::View, + "CREATE VIEW `active_users` AS SELECT `id` FROM `users`", + ); + + assert_eq!(sql, "CREATE VIEW `active_users` AS SELECT `id` FROM `users`;"); + } + #[test] fn sqlserver_view_source_opened_for_editing_shows_alter_view() { let sql = build_editable_object_source(EditableObjectSourceSqlInput {