From b6a39c1fa03073075a0c94586ad674ce55bc5ddf Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 6 May 2026 18:22:41 +0800 Subject: [PATCH] feat(query): set schema context (search_path) before query execution (#113) Add schema selector to the editor toolbar for schema-aware databases (PostgreSQL, SQL Server, Oracle). When a schema is selected, the backend sets the appropriate search_path before executing queries so users no longer need to write schema-qualified table names. --- crates/dbx-core/src/db/postgres.rs | 60 +++++++++++++++++++++++++ crates/dbx-core/src/query.rs | 23 +++++++--- src-tauri/src/commands/query.rs | 20 +++++++-- src-tauri/src/commands/sql_file.rs | 1 + src-web/src/routes/query.rs | 28 +++++++++--- src-web/src/routes/sql_file.rs | 2 +- src/App.vue | 6 +++ src/components/layout/EditorToolbar.vue | 57 ++++++++++++++++++++++- src/composables/useSchemaOptions.ts | 50 +++++++++++++++++++++ src/i18n/locales/en.ts | 1 + src/i18n/locales/zh-CN.ts | 1 + src/lib/http.ts | 24 +++++++--- src/lib/tauri.ts | 24 +++++++--- src/stores/queryStore.ts | 13 +++++- src/types/database.ts | 1 + 15 files changed, 279 insertions(+), 32 deletions(-) create mode 100644 src/composables/useSchemaOptions.ts diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index f52a0f621b..1443bdb102 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -300,6 +300,66 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result Result { + let mut conn = pool.acquire().await.map_err(|e| e.to_string())?; + let set_path = format!("SET search_path TO \"{}\", public", schema); + sqlx::query(&set_path).execute(&mut *conn).await.map_err(|e| e.to_string())?; + + let start = Instant::now(); + let trimmed = sql.trim().to_uppercase(); + + if trimmed.starts_with("SELECT") + || trimmed.starts_with("SHOW") + || trimmed.starts_with("EXPLAIN") + || trimmed.starts_with("WITH") + || trimmed.starts_with("TABLE") + { + let rows: Vec = + sqlx::query(sql).persistent(false).fetch_all(&mut *conn).await.map_err(|e| e.to_string())?; + + let (columns, column_types): (Vec, Vec) = if let Some(first) = rows.first() { + let cols = first.columns(); + ( + cols.iter().map(|c| c.name().to_string()).collect(), + cols.iter().map(|c| c.type_info().name().to_string()).collect(), + ) + } else { + let desc = (&mut *conn).describe(sql).await.map_err(|e| e.to_string())?; + ( + desc.columns().iter().map(|c| c.name().to_string()).collect(), + desc.columns().iter().map(|c| c.type_info().name().to_string()).collect(), + ) + }; + + let result_rows: Vec> = rows + .iter() + .map(|row| { + (0..row.len()) + .map(|i| pg_value_to_json(row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .collect() + }) + .collect(); + + Ok(QueryResult { + columns, + rows: result_rows, + affected_rows: 0, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + }) + } else { + let result = sqlx::query(sql).execute(&mut *conn).await.map_err(|e| e.to_string())?; + + Ok(QueryResult { + columns: vec![], + rows: vec![], + affected_rows: result.rows_affected(), + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + }) + } +} + pub async fn list_indexes(pool: &PgPool, schema: &str, table: &str) -> Result, String> { let rows: Vec = sqlx::query( "SELECT i.relname AS index_name, \ diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 5a2e4deb07..e548c7077e 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -134,6 +134,7 @@ pub async fn do_execute( state: &AppState, pool_key: &str, sql: &str, + schema: Option<&str>, cancel_token: Option, ) -> Result { let connections = state.connections.lock().await; @@ -161,8 +162,15 @@ pub async fn do_execute( } PoolKind::Postgres(p) => { let p = p.clone(); + let schema = schema.map(|s| s.to_string()); drop(connections); - wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)).await.map(truncate_result) + if let Some(schema) = schema { + wait_for_query(cancel_token, db::postgres::execute_query_with_schema(&p, &schema, sql)) + .await + .map(truncate_result) + } else { + wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)).await.map(truncate_result) + } } PoolKind::Sqlite(p) => { let p = p.clone(); @@ -214,6 +222,7 @@ pub async fn execute_sql_statement( connection_id: &str, database: &str, sql: &str, + schema: Option<&str>, cancel_token: Option, ) -> Result { let pool_key = if database.is_empty() { @@ -226,13 +235,13 @@ pub async fn execute_sql_statement( return Err(canceled_error()); } - let result = do_execute(state, &pool_key, sql, cancel_token.clone()).await; + let result = do_execute(state, &pool_key, sql, schema, cancel_token.clone()).await; match &result { Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; let new_key = state.reconnect_pool(connection_id, db_opt).await?; - do_execute(state, &new_key, sql, cancel_token).await + do_execute(state, &new_key, sql, schema, cancel_token).await } _ => result, } @@ -243,12 +252,13 @@ pub async fn execute_multi_core( connection_id: &str, database: &str, sql: &str, + schema: Option<&str>, cancel_token: Option, ) -> Result, String> { let statements = split_sql_statements(sql); if statements.len() <= 1 { let single_sql = statements.into_iter().next().unwrap_or_default(); - let result = execute_sql_statement(state, connection_id, database, &single_sql, cancel_token).await?; + let result = execute_sql_statement(state, connection_id, database, &single_sql, schema, cancel_token).await?; return Ok(vec![result]); } @@ -264,7 +274,7 @@ pub async fn execute_multi_core( }); break; } - match execute_sql_statement(state, connection_id, database, stmt, cancel_token.clone()).await { + match execute_sql_statement(state, connection_id, database, stmt, schema, cancel_token.clone()).await { Ok(r) => results.push(r), Err(e) => { results.push(db::QueryResult { @@ -286,6 +296,7 @@ pub async fn execute_statements( connection_id: &str, database: &str, statements: &[String], + schema: Option<&str>, ) -> Result { let pool_key = if database.is_empty() { connection_id.to_string() @@ -297,7 +308,7 @@ pub async fn execute_statements( let start = std::time::Instant::now(); for (i, sql) in statements.iter().enumerate() { - match do_execute(state, &pool_key, sql, None).await { + match do_execute(state, &pool_key, sql, schema, None).await { Ok(result) => { total_affected += result.affected_rows; } diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 8f7da1039a..eacdae4987 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -14,13 +14,15 @@ pub async fn execute_query( connection_id: String, database: String, sql: String, + schema: Option, execution_id: Option, ) -> Result { let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); let cancel_token = registered_query.as_ref().map(|query| query.token()); - dbx_core::query::execute_sql_statement(&state, &connection_id, &database, &sql, cancel_token).await + dbx_core::query::execute_sql_statement(&state, &connection_id, &database, &sql, schema.as_deref(), cancel_token) + .await } #[tauri::command] @@ -29,13 +31,14 @@ pub async fn execute_multi( connection_id: String, database: String, sql: String, + schema: Option, execution_id: Option, ) -> Result, String> { let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); let cancel_token = registered_query.as_ref().map(|query| query.token()); - dbx_core::query::execute_multi_core(&state, &connection_id, &database, &sql, cancel_token).await + dbx_core::query::execute_multi_core(&state, &connection_id, &database, &sql, schema.as_deref(), cancel_token).await } #[tauri::command] @@ -49,8 +52,9 @@ pub async fn execute_batch( connection_id: String, database: String, statements: Vec, + schema: Option, ) -> Result { - dbx_core::query::execute_statements(&state, &connection_id, &database, &statements).await + dbx_core::query::execute_statements(&state, &connection_id, &database, &statements, schema.as_deref()).await } #[tauri::command] @@ -59,6 +63,14 @@ pub async fn execute_script( connection_id: String, database: String, sql: String, + schema: Option, ) -> Result { - dbx_core::query::execute_statements(&state, &connection_id, &database, &split_sql_statements(&sql)).await + dbx_core::query::execute_statements( + &state, + &connection_id, + &database, + &split_sql_statements(&sql), + schema.as_deref(), + ) + .await } diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index 1c16546ff2..f73d7d696c 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -262,6 +262,7 @@ async fn execute_statement_with_progress( &request.connection_id, &request.database, statement, + None, Some(token.clone()), ) .await diff --git a/src-web/src/routes/query.rs b/src-web/src/routes/query.rs index dfcd77961d..efd2402625 100644 --- a/src-web/src/routes/query.rs +++ b/src-web/src/routes/query.rs @@ -13,6 +13,7 @@ pub struct ExecuteQueryRequest { pub connection_id: String, pub database: String, pub sql: String, + pub schema: Option, pub execution_id: Option, } @@ -28,6 +29,7 @@ pub struct ExecuteBatchRequest { pub connection_id: String, pub database: String, pub statements: Vec, + pub schema: Option, } pub async fn execute_query( @@ -44,6 +46,7 @@ pub async fn execute_query( &req.connection_id, &req.database, &req.sql, + req.schema.as_deref(), Some(cancel_token), ) .await @@ -67,6 +70,7 @@ pub async fn execute_multi( &req.connection_id, &req.database, &req.sql, + req.schema.as_deref(), Some(cancel_token), ) .await @@ -80,9 +84,15 @@ pub async fn execute_batch( State(state): State>, Json(req): Json, ) -> Result, AppError> { - let result = dbx_core::query::execute_statements(&state.app, &req.connection_id, &req.database, &req.statements) - .await - .map_err(AppError)?; + let result = dbx_core::query::execute_statements( + &state.app, + &req.connection_id, + &req.database, + &req.statements, + req.schema.as_deref(), + ) + .await + .map_err(AppError)?; Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } @@ -100,9 +110,15 @@ pub async fn execute_script( Json(req): Json, ) -> Result, AppError> { let statements = dbx_core::sql::split_sql_statements(&req.sql); - let result = dbx_core::query::execute_statements(&state.app, &req.connection_id, &req.database, &statements) - .await - .map_err(AppError)?; + let result = dbx_core::query::execute_statements( + &state.app, + &req.connection_id, + &req.database, + &statements, + req.schema.as_deref(), + ) + .await + .map_err(AppError)?; Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } diff --git a/src-web/src/routes/sql_file.rs b/src-web/src/routes/sql_file.rs index 6dbb5a69a6..a3a71a73e3 100644 --- a/src-web/src/routes/sql_file.rs +++ b/src-web/src/routes/sql_file.rs @@ -138,7 +138,7 @@ pub async fn execute_sql_file( let _ = tx.send(json); } - match query::execute_sql_statement(&app, &req.connection_id, &req.database, stmt, None).await { + match query::execute_sql_statement(&app, &req.connection_id, &req.database, stmt, None, None).await { Ok(result) => { success_count += 1; total_affected += result.affected_rows; diff --git a/src/App.vue b/src/App.vue index 3c03628c18..e5b66c7232 100644 --- a/src/App.vue +++ b/src/App.vue @@ -248,6 +248,11 @@ function changeActiveDatabase(database: string) { if (tab) queryStore.updateDatabase(tab.id, database); } +function changeActiveSchema(schema: string | undefined) { + const tab = activeTab.value; + if (tab) queryStore.updateSchema(tab.id, schema); +} + function toggleLocale() { setLocale(currentLocale() === "zh-CN" ? "en" : "zh-CN"); } @@ -426,6 +431,7 @@ onUnmounted(() => { @open-sql="openSqlFile" @change-connection="changeActiveConnection" @change-database="changeActiveDatabase" + @change-schema="changeActiveSchema" /> import { computed } from "vue"; import { useI18n } from "vue-i18n"; -import { Play, Loader2, Square, Database, Table2, AlignLeft, GitBranch, Save, FolderOpen } from "lucide-vue-next"; +import { + Play, + Loader2, + Square, + Database, + Table2, + AlignLeft, + GitBranch, + Save, + FolderOpen, + Layers, +} from "lucide-vue-next"; import { Button } from "@/components/ui/button"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Tooltip, TooltipTrigger, TooltipContent } from "@/components/ui/tooltip"; import DatabaseIcon from "@/components/icons/DatabaseIcon.vue"; import { useConnectionStore } from "@/stores/connectionStore"; import { useDatabaseOptions } from "@/composables/useDatabaseOptions"; +import { useSchemaOptions } from "@/composables/useSchemaOptions"; import { connectionIconType } from "@/lib/connectionPresentation"; import { connectionDisplayName } from "@/lib/tabPresentation"; import type { QueryTab, ConnectionConfig } from "@/types/database"; @@ -27,11 +39,13 @@ const emit = defineEmits<{ openSql: []; changeConnection: [connectionId: string]; changeDatabase: [database: string]; + changeSchema: [schema: string | undefined]; }>(); const { t } = useI18n(); const connectionStore = useConnectionStore(); const { databaseOptions, loadingDatabaseOptions, loadDatabaseOptions } = useDatabaseOptions(); +const { loadSchemaOptions, getSchemaOptionsForDb, isLoadingSchemas, isSchemaAware } = useSchemaOptions(); const activeDatabaseOptions = computed(() => { const connection = props.activeConnection; @@ -40,6 +54,18 @@ const activeDatabaseOptions = computed(() => { const activeDatabaseValue = computed(() => props.activeTab.database || ""); const activeConnectionValue = computed(() => props.activeConnection?.id || ""); +const activeSchemaValue = computed(() => props.activeTab.schema || ""); + +const showSchemaSelector = computed(() => { + const connection = props.activeConnection; + return connection && isSchemaAware(connection.id) && props.activeTab.database; +}); + +const activeSchemaOptions = computed(() => { + const connection = props.activeConnection; + if (!connection) return []; + return getSchemaOptionsForDb(connection.id, props.activeTab.database); +}); function databaseDisplayName(database: string): string { const connection = props.activeConnection; @@ -185,6 +211,35 @@ function databaseDisplayName(database: string): string { +
+ + +
diff --git a/src/composables/useSchemaOptions.ts b/src/composables/useSchemaOptions.ts new file mode 100644 index 0000000000..dda19d7421 --- /dev/null +++ b/src/composables/useSchemaOptions.ts @@ -0,0 +1,50 @@ +import { ref } from "vue"; +import { useConnectionStore } from "@/stores/connectionStore"; +import * as api from "@/lib/api"; + +export function useSchemaOptions() { + const connectionStore = useConnectionStore(); + + const schemaOptions = ref>({}); + const loadingSchemaOptions = ref>({}); + + function cacheKey(connectionId: string, database: string) { + return `${connectionId}:${database}`; + } + + function isSchemaAware(connectionId: string): boolean { + const dbType = connectionStore.getConfig(connectionId)?.db_type; + return dbType === "postgres" || dbType === "sqlserver" || dbType === "oracle"; + } + + async function loadSchemaOptions(connectionId: string, database: string) { + if (!isSchemaAware(connectionId) || !database) return; + const key = cacheKey(connectionId, database); + if (loadingSchemaOptions.value[key]) return; + + loadingSchemaOptions.value[key] = true; + try { + await connectionStore.ensureConnected(connectionId); + schemaOptions.value[key] = await api.listSchemas(connectionId, database); + } finally { + loadingSchemaOptions.value[key] = false; + } + } + + function getSchemaOptionsForDb(connectionId: string, database: string): string[] { + return schemaOptions.value[cacheKey(connectionId, database)] ?? []; + } + + function isLoadingSchemas(connectionId: string, database: string): boolean { + return !!loadingSchemaOptions.value[cacheKey(connectionId, database)]; + } + + return { + schemaOptions, + loadingSchemaOptions, + loadSchemaOptions, + getSchemaOptionsForDb, + isLoadingSchemas, + isSchemaAware, + }; +} diff --git a/src/i18n/locales/en.ts b/src/i18n/locales/en.ts index aae68e7528..19e58d9f12 100644 --- a/src/i18n/locales/en.ts +++ b/src/i18n/locales/en.ts @@ -128,6 +128,7 @@ export default { noDatabase: "No database selected", selectConnection: "Select connection", selectDatabase: "Select database", + selectSchema: "Select schema", }, tabs: { sql: "SQL", diff --git a/src/i18n/locales/zh-CN.ts b/src/i18n/locales/zh-CN.ts index a2885f03c8..b6e2707d24 100644 --- a/src/i18n/locales/zh-CN.ts +++ b/src/i18n/locales/zh-CN.ts @@ -127,6 +127,7 @@ export default { noDatabase: "未选择数据库", selectConnection: "选择连接", selectDatabase: "选择数据库", + selectSchema: "选择模式", }, tabs: { sql: "SQL", diff --git a/src/lib/http.ts b/src/lib/http.ts index dc570df509..e179df59a7 100644 --- a/src/lib/http.ts +++ b/src/lib/http.ts @@ -157,26 +157,38 @@ export async function executeQuery( connectionId: string, database: string, sql: string, + schema?: string, executionId?: string, ): Promise { - return post("/api/query/execute", { connectionId, database, sql, executionId }); + return post("/api/query/execute", { connectionId, database, sql, schema, executionId }); } export async function executeMulti( connectionId: string, database: string, sql: string, + schema?: string, executionId?: string, ): Promise { - return post("/api/query/execute-multi", { connectionId, database, sql, executionId }); + return post("/api/query/execute-multi", { connectionId, database, sql, schema, executionId }); } -export async function executeBatch(connectionId: string, database: string, statements: string[]): Promise { - return post("/api/query/execute-batch", { connectionId, database, statements }); +export async function executeBatch( + connectionId: string, + database: string, + statements: string[], + schema?: string, +): Promise { + return post("/api/query/execute-batch", { connectionId, database, statements, schema }); } -export async function executeScript(connectionId: string, database: string, sql: string): Promise { - return post("/api/query/execute-script", { connectionId, database, sql }); +export async function executeScript( + connectionId: string, + database: string, + sql: string, + schema?: string, +): Promise { + return post("/api/query/execute-script", { connectionId, database, sql, schema }); } export async function cancelQuery(executionId: string): Promise { diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 1f959c079f..168f5397b3 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -138,30 +138,42 @@ export async function executeQuery( connectionId: string, database: string, sql: string, + schema?: string, executionId?: string, ): Promise { - return invoke("execute_query", { connectionId, database, sql, executionId }); + return invoke("execute_query", { connectionId, database, sql, schema, executionId }); } export async function executeMulti( connectionId: string, database: string, sql: string, + schema?: string, executionId?: string, ): Promise { - return invoke("execute_multi", { connectionId, database, sql, executionId }); + return invoke("execute_multi", { connectionId, database, sql, schema, executionId }); } export async function cancelQuery(executionId: string): Promise { return invoke("cancel_query", { executionId }); } -export async function executeBatch(connectionId: string, database: string, statements: string[]): Promise { - return invoke("execute_batch", { connectionId, database, statements }); +export async function executeBatch( + connectionId: string, + database: string, + statements: string[], + schema?: string, +): Promise { + return invoke("execute_batch", { connectionId, database, statements, schema }); } -export async function executeScript(connectionId: string, database: string, sql: string): Promise { - return invoke("execute_script", { connectionId, database, sql }); +export async function executeScript( + connectionId: string, + database: string, + sql: string, + schema?: string, +): Promise { + return invoke("execute_script", { connectionId, database, sql, schema }); } export async function listIndexes( diff --git a/src/stores/queryStore.ts b/src/stores/queryStore.ts index c0f20636b9..d07591f970 100644 --- a/src/stores/queryStore.ts +++ b/src/stores/queryStore.ts @@ -148,17 +148,25 @@ export const useQueryStore = defineStore("query", () => { const tab = tabs.value.find((t) => t.id === id); if (!tab || tab.database === database) return; tab.database = database; + tab.schema = undefined; tab.result = undefined; tab.lastExecutedSql = undefined; clearExplain(tab); tab.tableMeta = undefined; } + function updateSchema(id: string, schema: string | undefined) { + const tab = tabs.value.find((t) => t.id === id); + if (!tab || tab.schema === schema) return; + tab.schema = schema; + } + function updateConnection(id: string, connectionId: string, database = "") { const tab = tabs.value.find((t) => t.id === id); if (!tab || tab.connectionId === connectionId) return; tab.connectionId = connectionId; tab.database = database; + tab.schema = undefined; tab.result = undefined; tab.lastExecutedSql = undefined; clearExplain(tab); @@ -229,7 +237,7 @@ export const useQueryStore = defineStore("query", () => { tab.executionId = executionId; tab.lastExecutedSql = sql; try { - const results = await api.executeMulti(tab.connectionId, tab.database, sql, executionId); + const results = await api.executeMulti(tab.connectionId, tab.database, sql, tab.schema, executionId); const current = tabs.value.find((t) => t.id === id); if (current?.executionId === executionId) { if (results.length > 1) { @@ -278,7 +286,7 @@ export const useQueryStore = defineStore("query", () => { tab.explainSql = built.sql; tab.lastExplainedSql = sql; try { - const result = await api.executeQuery(tab.connectionId, tab.database, built.sql, executionId); + const result = await api.executeQuery(tab.connectionId, tab.database, built.sql, tab.schema, executionId); const current = tabs.value.find((t) => t.id === id); if (current?.explainExecutionId === executionId) { current.explainPlan = parseExplainResult(databaseType as "mysql" | "postgres", result); @@ -373,6 +381,7 @@ export const useQueryStore = defineStore("query", () => { updateSql, togglePinnedTab, updateDatabase, + updateSchema, updateConnection, setTableMeta, setExecuting, diff --git a/src/types/database.ts b/src/types/database.ts index 3225bf0bb4..5d49a9d219 100644 --- a/src/types/database.ts +++ b/src/types/database.ts @@ -147,6 +147,7 @@ export interface QueryTab { title: string; connectionId: string; database: string; + schema?: string; sql: string; lastExecutedSql?: string; pinned?: boolean;