feat(ai): cross-database query with multi database selection

This commit is contained in:
二丫讲梵
2026-09-14 13:50:17 +08:00
committed by GitHub
parent 58d902c07f
commit 193d875b8a
20 changed files with 240 additions and 57 deletions
@@ -50,7 +50,6 @@ import {
} from "@lucide/vue";
import { Button } from "@/components/ui/button";
import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from "@/components/ui/dropdown-menu";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Popover, PopoverAnchor, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
import { ScrollArea } from "@/components/ui/scroll-area";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
@@ -156,7 +155,7 @@ import type { AiConfigItem, AiEffortCapability, AiEffortOption, AiEffortSelectio
import type { ConnectionConfig, QueryTab, SavedSqlFile, TableInfo } from "@/types/database";
import { fetchNamespaceOptionsForConnection, useDatabaseOptions } from "@/composables/useDatabaseOptions";
import { useSchemaOptions } from "@/composables/useSchemaOptions";
import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase";
import { encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase";
import { normalizeSqliteNamespace } from "@/lib/database/sqliteNamespace";
import { isQueryExecutionErrorResult } from "@/lib/query/queryResultError";
import { isSchemaAware, isSingleDatabase } from "@/lib/database/databaseCapabilities";
@@ -1405,19 +1404,43 @@ const dbSelectOptions = computed(() => {
}));
});
const selectedNamespace = computed(() => (props.connection && props.tab ? resolveAiNamespaceSelection(props.tab, props.connection).value : ""));
const selectedDatabaseSelectValue = computed(() => (props.connection ? encodeSelectableDatabaseValue(props.connection.db_type, selectedNamespace.value) : ""));
// AI can inspect more than the tab's active database. Keep this selection local
// to the composer so changing the visible query tab does not rewrite SQL state.
const selectedDatabases = ref<string[]>([]);
const selectedDatabaseValues = computed(() => new Set(selectedDatabases.value));
const selectedDatabaseLabel = computed(() => {
if (!props.connection) return t("editor.selectDatabase");
if (!props.tab) return t("editor.selectDatabase");
return formatDatabaseLabel(props.connection, selectedNamespace.value, {
defaultDatabase: t("editor.defaultDatabase"),
noDatabase: t("editor.noDatabase"),
});
const labels = dbSelectOptions.value.filter((option) => selectedDatabaseValues.value.has(option.database)).map((option) => option.label);
return labels.length ? labels.join(", ") : t("editor.selectDatabase");
});
function syncSelectedDatabases() {
const active = selectedNamespace.value;
const available = dbSelectOptions.value.map((option) => option.database);
const retained = selectedDatabases.value.filter((database) => available.includes(database));
selectedDatabases.value = retained.length ? retained : active ? [active] : [];
}
function toggleDatabase(database: string) {
if (selectedDatabaseValues.value.has(database)) {
if (selectedDatabases.value.length === 1) return;
selectedDatabases.value = selectedDatabases.value.filter((item) => item !== database);
} else {
selectedDatabases.value = [...selectedDatabases.value, database];
}
}
const selectedNamespace = computed(() => (props.connection && props.tab ? resolveAiNamespaceSelection(props.tab, props.connection).value : ""));
watch(
() => `${props.connection?.id ?? ""}:${props.tab?.id ?? ""}`,
() => {
selectedDatabases.value = selectedNamespace.value ? [selectedNamespace.value] : [];
},
);
watch([dbSelectOptions, selectedNamespace], syncSelectedDatabases, { immediate: true });
const showAiSchemaSelector = computed(() => {
const connection = props.connection;
return !!connection && connection.db_type !== "dameng" && aiSchemaSelectionSupported(connection);
@@ -1478,20 +1501,6 @@ async function changeConnection(connectionId: string) {
}
}
function changeNamespace(value: string) {
const tab = props.tab;
const connection = props.connection;
if (!tab || !connection) return;
const namespace = decodeSelectableDatabaseValue(connection.db_type, value);
if (resolveAiNamespaceSelection(tab, connection).value === namespace) return;
clearContextReferences();
if (resolveAiNamespaceSelection(tab, connection).kind === "schema") {
queryStore.updateSchema(tab.id, namespace || undefined);
} else {
queryStore.updateDatabase(tab.id, namespace);
}
}
function changeSchema(schema: string) {
const tab = props.tab;
if (!tab || tab.schema === schema) return;
@@ -2793,6 +2802,9 @@ async function send() {
clearPendingWriteGrant();
return;
}
// Capture the selection before context loading or queued run scheduling can
// yield to another conversation. Dameng's top-level selector is a schema.
const runDatabases = resolveAiNamespaceSelection(tab, connection).kind === "database" ? [...selectedDatabases.value] : [];
const activeConfig = activeFullConfig.value;
if (!activeConfig) {
clearPendingWriteGrant();
@@ -3086,11 +3098,21 @@ async function send() {
// paying for buildAiContext() too; it can do real backend/schema work that
// would be entirely wasted on an already-abandoned request.
if (!generationCanContinue()) return;
const context = await buildAiContext(tab, connection, {
mentionedTables,
sqlFiles,
csvFiles: csvAttachments,
});
const requestDatabase = runDatabases[0] ?? tab.database;
const context = await buildAiContext(
{
...tab,
database: requestDatabase,
schema: runDatabases.length > 1 || requestDatabase !== tab.database ? undefined : tab.schema,
},
connection,
{
mentionedTables,
sqlFiles,
csvFiles: csvAttachments,
},
);
context.selectedDatabases = runDatabases;
// Superseded while awaiting buildAiContext() above — must bail before ever
// calling runAgentStream(), not just before writing its results. Without
// this recheck, a clear/switch/unmount that fires during context
@@ -4970,27 +4992,26 @@ async function openExternalUrl(url: string) {
/>
<template v-if="connection">
<Database class="h-3 w-3 shrink-0 text-foreground/40" />
<Select
:model-value="selectedDatabaseSelectValue"
@update:model-value="
(v) => {
if (typeof v === 'string') changeNamespace(v);
}
"
<Popover
@update:open="
(open: boolean) => {
if (open) loadDatabases();
}
"
>
<SelectTrigger :class="['h-5 w-auto border-0 rounded-md bg-transparent dark:bg-transparent p-0 px-1 text-xs text-foreground/80 shadow-none focus:ring-0 focus-visible:ring-0 [&_svg]:size-3', showAiSchemaSelector && 'min-w-0 max-w-56 flex-1']">
<SelectValue :placeholder="t('editor.selectDatabase')">{{ selectedDatabaseLabel }}</SelectValue>
</SelectTrigger>
<SelectContent>
<SelectItem v-for="option in dbSelectOptions" :key="option.value" :value="option.value">{{ option.label }}</SelectItem>
<SelectItem v-if="!dbSelectOptions.length && connection && tab" :value="selectedDatabaseSelectValue">{{ selectedDatabaseLabel }}</SelectItem>
</SelectContent>
</Select>
<PopoverTrigger as-child>
<Button variant="ghost" :class="['h-5 max-w-64 justify-start border-0 p-0 px-1 text-xs font-normal text-foreground/80 shadow-none', showAiSchemaSelector && 'min-w-0 flex-1']">
<span class="truncate">{{ selectedDatabaseLabel }}</span>
</Button>
</PopoverTrigger>
<PopoverContent align="start" class="w-64 p-1">
<button v-for="option in dbSelectOptions" :key="option.value" type="button" class="flex w-full items-center gap-2 rounded-sm px-2 py-1.5 text-left text-sm hover:bg-muted" @click="toggleDatabase(option.database)">
<Check :class="['h-4 w-4', selectedDatabaseValues.has(option.database) ? 'opacity-100' : 'opacity-0']" />
<span class="truncate">{{ option.label }}</span>
</button>
<div v-if="!dbSelectOptions.length" class="px-2 py-1.5 text-sm text-muted-foreground">{{ t("editor.selectDatabase") }}</div>
</PopoverContent>
</Popover>
<template v-if="showAiSchemaSelector">
<Layers class="h-3 w-3 shrink-0 text-foreground/40" />
<SearchableSelect
+28 -2
View File
@@ -1,5 +1,6 @@
import { beforeAll, describe, expect, it } from "vitest";
import { buildAgentRequest, buildSystemPrompt, buildUserPrompt, type AiContext } from "@/lib/ai/ai";
import { beforeAll, describe, expect, it, vi } from "vitest";
import { buildAgentRequest, buildSystemPrompt, buildUserPrompt, runAgentStream, type AiContext } from "@/lib/ai/ai";
import * as api from "@/lib/backend/api";
import { setLocale } from "@/i18n";
function context(overrides: Partial<AiContext> = {}): AiContext {
@@ -18,6 +19,31 @@ function context(overrides: Partial<AiContext> = {}): AiContext {
}
describe("AI SQL dialect prompt", () => {
it("sends database selections to the backend as well as the model", async () => {
const stream = vi.spyOn(api, "aiAgentStream").mockResolvedValue("done");
try {
await runAgentStream(
{
config: { provider: "openai", apiKey: "test", apiUrl: "https://example.invalid", model: "model" },
action: "general",
mode: "agent",
instruction: "Join users and orders",
context: context({ databaseType: "mysql", database: "db_a", selectedDatabases: ["db_a", "db_b"] }),
},
[],
() => {},
"multi-db-run",
);
const args = stream.mock.calls[0];
expect(args[3]).toBe("db_a");
expect(args[14]).toEqual(["db_a", "db_b"]);
expect(args[1].systemPrompt).toContain('Selected databases: ["db_a","db_b"]');
expect(args[8]).toBe(false);
} finally {
stream.mockRestore();
}
});
// buildSystemPrompt picks zh/en copy via currentLocale(); pin to en so the
// English-string assertions are deterministic regardless of the host OS locale.
beforeAll(async () => {
+14 -2
View File
@@ -100,6 +100,8 @@ export interface AiContext {
connectionName: string;
databaseType: DatabaseType;
database: string;
/** Databases selected for this run; omitted by older callers. */
selectedDatabases?: string[];
/** Schema selected for metadata loading and agent tool execution. */
schema?: string;
currentSql: string;
@@ -221,7 +223,8 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa
const { messages, systemPrompt, taskContract, maxTokens } = buildAgentRequest(input, history, custom);
const sid = sessionId || uuid();
return api.aiAgentStream(
const selectedDatabases = input.context.selectedDatabases;
const args = [
sid,
{
config: input.config,
@@ -241,7 +244,11 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa
input.confirmedConnectionId,
input.confirmedDatabase,
input.confirmedSchema,
);
] as const;
if (selectedDatabases?.length) {
return api.aiAgentStream(...args, undefined, selectedDatabases);
}
return api.aiAgentStream(...args);
}
export function buildUserPrompt(action: AiAction, context: AiContext, instruction: string, isZh: boolean): string {
@@ -319,6 +326,11 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai
`Database type: ${context.databaseType}`,
`Connection: ${context.connectionName}`,
`Database: ${context.database}`,
context.selectedDatabases?.length
? isZh
? `已选择数据库:${JSON.stringify(context.selectedDatabases)}。请在元数据工具中通过 database 参数分别检查这些数据库。跨库 JOIN 需要数据库引擎支持:MySQL 使用 database.table,SQL Server 使用 database.schema.table;PostgreSQL 不能直接连接不同数据库,除非已配置联邦查询。MCP 授权仍然生效,选择数据库不会绕过授权限制。`
: `Selected databases: ${JSON.stringify(context.selectedDatabases)}. Use the database parameter on metadata tools to inspect each selected database. Cross-database joins require engine support: MySQL uses database.table, SQL Server uses database.schema.table; PostgreSQL cannot directly join separate databases without an existing federation setup. MCP authorization still applies; selecting databases does not override it.`
: "",
context.schema ? `Selected schema: ${context.schema}` : "",
schemaCoverageLine(context, isZh),
"",
+2
View File
@@ -1873,6 +1873,7 @@ export async function aiAgentStream(
confirmedDatabase?: string,
confirmedSchema?: string,
signal?: AbortSignal,
selectedDatabases?: string[],
): Promise<string> {
const res = await fetch(apiUrl("/api/ai/agent-stream"), {
method: "POST",
@@ -1890,6 +1891,7 @@ export async function aiAgentStream(
confirmedConnectionId,
confirmedDatabase,
confirmedSchema,
selectedDatabases,
}),
signal,
});
+2
View File
@@ -587,6 +587,7 @@ export async function aiAgentStream(
confirmedDatabase?: string,
confirmedSchema?: string,
_signal?: AbortSignal,
selectedDatabases?: string[],
): Promise<string> {
const unlisten: UnlistenFn = await listen<TauriAgentEvent>("ai-agent-event", (event) => {
const payload = event.payload;
@@ -610,6 +611,7 @@ export async function aiAgentStream(
confirmedConnectionId,
confirmedDatabase,
confirmedSchema,
selectedDatabases,
});
} catch (e) {
unlisten();
+2
View File
@@ -76,6 +76,7 @@ pub struct AgentLoopContext {
pub state: Arc<AppState>,
pub connection_id: String,
pub database: String,
pub selected_databases: Vec<String>,
/// Selected schema that scopes Agent metadata and SQL execution.
pub schema: Option<String>,
pub db_type: DatabaseType,
@@ -151,6 +152,7 @@ pub async fn run_agent_loop(
connection_id: agent_ctx.connection_id.clone(),
connection_name,
database: agent_ctx.database.clone(),
selected_databases: agent_ctx.selected_databases.clone(),
schema: agent_ctx.schema.clone(),
agent_mode: is_agent_mode,
allow_writes: agent_ctx.sql_permissions.allow_writes,
+59 -7
View File
@@ -55,7 +55,8 @@ fn connection_tool_lock(connection_id: &str) -> Arc<tokio::sync::Mutex<()>> {
fn tool_uses_database(tool_name: &str) -> bool {
matches!(
tool_name,
"list_tables"
"list_databases"
| "list_tables"
| "get_columns"
| "execute_query"
| "get_sample_data"
@@ -320,7 +321,7 @@ pub fn read_only_tools(db_type: DatabaseType) -> Vec<ToolDefinition> {
if is_vector_db(db_type) {
vec![list_collections_tool(), get_current_time_tool()]
} else {
vec![list_tables_tool(), get_columns_tool(), get_current_time_tool()]
vec![list_databases_tool(), list_tables_tool(), get_columns_tool(), get_current_time_tool()]
}
}
@@ -331,7 +332,7 @@ pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) ->
if is_vector_db(db_type) {
return vec![list_collections_tool(), browse_collection_tool(), get_current_time_tool()];
}
let mut tools = vec![list_tables_tool(), get_columns_tool(), get_current_time_tool()];
let mut tools = vec![list_databases_tool(), list_tables_tool(), get_columns_tool(), get_current_time_tool()];
if db_type == DatabaseType::MongoDb {
tools.push(mongo_execute_query_tool(sql_permissions));
} else if supports_sql_query(db_type) {
@@ -344,6 +345,16 @@ pub fn all_tools(db_type: DatabaseType, sql_permissions: AgentSqlPermissions) ->
tools
}
fn list_databases_tool() -> ToolDefinition {
ToolDefinition {
name: "list_databases",
description: "List databases available through the current connection. If more than one database is returned, cross-database read queries can use fully qualified names such as database.table (or database.schema.table for SQL Server).",
parameters: json!({"type": "object", "properties": {}, "required": []}),
read_only: true,
parallel_ok: true,
}
}
fn mongo_execute_query_tool(_sql_permissions: AgentSqlPermissions) -> ToolDefinition {
ToolDefinition {
name: "execute_query",
@@ -378,6 +389,10 @@ fn list_tables_tool() -> ToolDefinition {
"schema": {
"type": "string",
"description": "Schema name to list tables from (optional, defaults to current database)"
},
"database": {
"type": "string",
"description": "Database to inspect (optional, defaults to the current database)"
}
},
"required": []
@@ -405,6 +420,10 @@ fn get_columns_tool() -> ToolDefinition {
"schema": {
"type": "string",
"description": "Schema name (optional, defaults to current database)"
},
"database": {
"type": "string",
"description": "Database containing the table (optional, defaults to the current database)"
}
},
"required": ["table"]
@@ -422,7 +441,7 @@ fn execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefinition {
} else if sql_permissions.allow_writes {
"Execute SQL after the user explicitly confirmed this operation. Read queries and non-DDL writes are allowed for this run."
} else {
"Execute a read-only SQL query and return results (max 50 rows). This run cannot execute writes or DDL because no specific SQL has been confirmed yet; this does not mean the database itself is read-only. When the user requests a write, first propose the exact SQL in one ```sql code block and ask for confirmation. After confirmation, DBX starts a new run that can execute only that exact SQL. Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements may be executed in this run."
"Execute a read-only SQL query and return results (max 50 rows). Cross-database reads may use fully qualified names such as database.table (or database.schema.table for SQL Server) without switching the current database. This run cannot execute writes or DDL because no specific SQL has been confirmed yet; this does not mean the database itself is read-only. When the user requests a write, first propose the exact SQL in one ```sql code block and ask for confirmation. After confirmation, DBX starts a new run that can execute only that exact SQL. Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements may be executed in this run."
};
ToolDefinition {
name: "execute_query",
@@ -478,6 +497,10 @@ fn get_sample_data_tool() -> ToolDefinition {
"type": "string",
"description": "Schema name (optional)"
},
"database": {
"type": "string",
"description": "Database containing the table (optional, defaults to the current database)"
},
"limit": {
"type": "number",
"description": "Max rows (default 20)"
@@ -572,6 +595,7 @@ pub async fn execute_tool(
None
};
let result = match tool_call.name.as_str() {
"list_databases" => execute_list_databases(tool_call, state, connection_id).await,
"list_tables" => execute_list_tables(tool_call, state, connection_id, database, default_schema, db_type).await,
"get_columns" => execute_get_columns(tool_call, state, connection_id, database, default_schema, db_type).await,
"execute_query" => {
@@ -629,6 +653,20 @@ pub async fn execute_tool(
}
}
async fn execute_list_databases(
_tool_call: &ToolCall,
state: &Arc<AppState>,
connection_id: &str,
) -> Result<String, String> {
let databases = crate::schema::list_databases_core(state, connection_id)
.await
.map_err(|e| format!("Failed to list databases: {e}"))?;
if databases.is_empty() {
return Ok("No databases were returned for this connection.".to_string());
}
Ok(databases.into_iter().map(|database| format!("- {}", database.name)).collect::<Vec<_>>().join("\n"))
}
async fn execute_list_tables(
tool_call: &ToolCall,
state: &Arc<AppState>,
@@ -637,13 +675,14 @@ async fn execute_list_tables(
default_schema: Option<&str>,
_db_type: &DatabaseType,
) -> Result<String, String> {
let database = effective_database(tool_call, database);
let schema = effective_schema(tool_call, default_schema).unwrap_or_default();
// Request one extra to detect whether more tables exist beyond the limit.
let tables = crate::schema::list_tables_core(
state,
connection_id,
database,
&database,
&schema,
None,
Some(LIST_TABLES_LIMIT + 1),
@@ -689,6 +728,7 @@ async fn execute_get_columns(
default_schema: Option<&str>,
_db_type: &DatabaseType,
) -> Result<String, String> {
let database = effective_database(tool_call, database);
let table = tool_call
.arguments
.get("table")
@@ -710,7 +750,7 @@ async fn execute_get_columns(
let schema = effective_schema(tool_call, default_schema).unwrap_or_default();
let columns = crate::schema::get_columns_core(state, connection_id, database, &schema, &table)
let columns = crate::schema::get_columns_core(state, connection_id, &database, &schema, &table)
.await
.map_err(|e| format!("Failed to get columns for {table}: {e}"))?;
@@ -999,6 +1039,7 @@ async fn execute_get_sample_data(
default_schema: Option<&str>,
db_type: &DatabaseType,
) -> Result<String, String> {
let database = effective_database(tool_call, database);
let table =
tool_call.arguments.get("table").and_then(|v| v.as_str()).ok_or("Missing required parameter: table")?.trim();
@@ -1032,7 +1073,7 @@ async fn execute_get_sample_data(
&synthetic_call,
state,
connection_id,
database,
&database,
schema.as_deref(),
db_type,
AgentSqlPermissions::default(),
@@ -1171,6 +1212,17 @@ fn effective_schema(tool_call: &ToolCall, default_schema: Option<&str>) -> Optio
})
}
fn effective_database(tool_call: &ToolCall, default_database: &str) -> String {
tool_call
.arguments
.get("database")
.and_then(Value::as_str)
.map(str::trim)
.filter(|database| !database.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| default_database.to_string())
}
/// Execute list_collections tool (vector databases).
async fn execute_list_collections(
_tool_call: &ToolCall,
@@ -602,6 +602,7 @@ mod tests {
connection_id: "conn-1".to_string(),
connection_name: "local".to_string(),
database: "demo".to_string(),
selected_databases: Vec::new(),
schema: None,
agent_mode: true,
allow_writes: false,
+47 -2
View File
@@ -14,6 +14,7 @@ pub struct CliAgentRunOptions {
pub connection_id: String,
pub connection_name: String,
pub database: String,
pub selected_databases: Vec<String>,
pub schema: Option<String>,
pub agent_mode: bool,
pub allow_writes: bool,
@@ -72,14 +73,27 @@ pub fn dbx_mcp_enabled_tools(agent_mode: bool) -> Vec<&'static str> {
}
pub fn dbx_mcp_scope_env(options: &CliAgentRunOptions) -> Vec<(&'static str, String)> {
let mut databases: Vec<&str> = Vec::new();
for database in &options.selected_databases {
let database = database.trim();
if !database.is_empty() && !databases.contains(&database) {
databases.push(database);
}
}
let multiple_databases = databases.len() > 1;
// An empty value explicitly clears inherited/configured single-database
// bounds. MCP's persisted connection/database policy still applies.
let database_scope = if multiple_databases { "" } else { databases.first().copied().unwrap_or(&options.database) };
let mut env = vec![
("DBX_MCP_ALLOW_WRITES", if options.allow_writes { "1" } else { "0" }.to_string()),
("DBX_MCP_ALLOW_DANGEROUS_SQL", if options.allow_dangerous { "1" } else { "0" }.to_string()),
("DBX_MCP_SCOPE_CONNECTION_ID", options.connection_id.clone()),
("DBX_MCP_SCOPE_CONNECTION_NAME", options.connection_name.clone()),
("DBX_MCP_SCOPE_DATABASE", options.database.clone()),
("DBX_MCP_SCOPE_DATABASE", database_scope.to_string()),
];
if let Some(schema) = options.schema.as_deref().filter(|schema| !schema.trim().is_empty()) {
if multiple_databases {
env.push(("DBX_MCP_SCOPE_SCHEMA", String::new()));
} else if let Some(schema) = options.schema.as_deref().filter(|schema| !schema.trim().is_empty()) {
env.push(("DBX_MCP_SCOPE_SCHEMA", schema.to_string()));
}
if let Some(ref sql) = options.confirmed_write_sql {
@@ -92,12 +106,43 @@ pub fn dbx_mcp_scope_env(options: &CliAgentRunOptions) -> Vec<(&'static str, Str
mod scope_env_tests {
use super::*;
#[test]
fn database_selection_controls_runtime_bounds_without_changing_write_permissions() {
for (selection, expected_database, expected_schema) in [
(vec![], "db_a", "public"),
(vec!["db_b"], "db_b", "public"),
(vec!["db_a", "db_b"], "", ""),
(vec![" db_a ", "db_a", " "], "db_a", "public"),
] {
let options = CliAgentRunOptions {
connection_id: "conn-1".into(),
connection_name: "Test".into(),
database: "db_a".into(),
selected_databases: selection.into_iter().map(str::to_string).collect(),
schema: Some("public".into()),
agent_mode: true,
allow_writes: false,
allow_dangerous: false,
confirmed_write_sql: None,
mcp_server_command: None,
};
let env: std::collections::HashMap<_, _> = dbx_mcp_scope_env(&options).into_iter().collect();
assert_eq!(env["DBX_MCP_SCOPE_DATABASE"], expected_database);
assert_eq!(env["DBX_MCP_SCOPE_SCHEMA"], expected_schema);
assert_eq!(env["DBX_MCP_SCOPE_CONNECTION_ID"], "conn-1");
assert_eq!(env["DBX_MCP_ALLOW_WRITES"], "0");
assert_eq!(env["DBX_MCP_ALLOW_DANGEROUS_SQL"], "0");
assert!(!env.contains_key("DBX_MCP_CONFIRMED_WRITE_SQL"));
}
}
#[test]
fn confirmed_write_sql_is_passed_to_the_scoped_mcp_subprocess() {
let options = CliAgentRunOptions {
connection_id: "dameng-1".to_string(),
connection_name: "Dameng".to_string(),
database: "APPDB".to_string(),
selected_databases: Vec::new(),
schema: Some("REPORTING".to_string()),
agent_mode: true,
allow_writes: true,
+3
View File
@@ -957,6 +957,7 @@ mod tests {
connection_id: "conn-1".to_string(),
connection_name: "local".to_string(),
database: "demo".to_string(),
selected_databases: Vec::new(),
schema: None,
agent_mode: true,
allow_writes: false,
@@ -1385,6 +1386,7 @@ exit 9
connection_id: String::new(),
connection_name: connection_name.clone(),
database,
selected_databases: Vec::new(),
schema: None,
agent_mode: false,
allow_writes: false,
@@ -1440,6 +1442,7 @@ exit 9
connection_id: String::new(),
connection_name,
database,
selected_databases: Vec::new(),
schema: None,
agent_mode: true,
allow_writes: false,
+1
View File
@@ -1148,6 +1148,7 @@ mod tests {
connection_id: "conn-1".to_string(),
connection_name: "local".to_string(),
database: "demo".to_string(),
selected_databases: Vec::new(),
schema: None,
agent_mode: true,
allow_writes: false,
+2
View File
@@ -517,6 +517,7 @@ mod tests {
connection_id: "sqlite-1".to_string(),
connection_name: "SQLite".to_string(),
database: "main".to_string(),
selected_databases: Vec::new(),
schema: None,
agent_mode: false,
allow_writes: false,
@@ -571,6 +572,7 @@ mod tests {
connection_id: String::new(),
connection_name: connection_name.clone(),
database,
selected_databases: Vec::new(),
schema: None,
agent_mode: false,
allow_writes: false,
+1
View File
@@ -819,6 +819,7 @@ mod tests {
connection_id: "conn-1".to_string(),
connection_name: "Demo".to_string(),
database: "app".to_string(),
selected_databases: Vec::new(),
schema: None,
agent_mode: true,
allow_writes: false,
+2
View File
@@ -485,6 +485,7 @@ mod tests {
connection_id: "conn-1".to_string(),
connection_name: "Local".to_string(),
database: "demo".to_string(),
selected_databases: Vec::new(),
schema: Some("public".to_string()),
agent_mode: true,
allow_writes: false,
@@ -654,6 +655,7 @@ mod tests {
connection_id: String::new(),
connection_name: connection_name.clone(),
database,
selected_databases: Vec::new(),
schema: None,
agent_mode: false,
allow_writes: false,
+1
View File
@@ -1017,6 +1017,7 @@ mod tests {
connection_id: "connection-1".to_string(),
connection_name: "Test connection".to_string(),
database: "dbx_test".to_string(),
selected_databases: Vec::new(),
schema: Some("reporting".to_string()),
agent_mode: true,
allow_writes: true,
+1
View File
@@ -561,6 +561,7 @@ mod tests {
connection_id: "conn-1".to_string(),
connection_name: "local".to_string(),
database: "demo".to_string(),
selected_databases: Vec::new(),
schema: Some("public".to_string()),
agent_mode: true,
allow_writes: false,
@@ -1,3 +1,5 @@
#![recursion_limit = "256"]
use std::sync::{Arc, Mutex};
use std::time::Duration;
+2
View File
@@ -1,3 +1,5 @@
#![recursion_limit = "256"]
pub mod backend;
pub mod diagnostics;
pub mod http;
+3
View File
@@ -93,6 +93,8 @@ pub struct AiAgentStreamRequest {
pub connection_id: String,
pub database: String,
#[serde(default)]
pub selected_databases: Vec<String>,
#[serde(default)]
pub schema: Option<String>,
pub db_type: String,
/// Agent mode: "ask" (read-only tools) or "agent" (all tools including execute_query).
@@ -452,6 +454,7 @@ pub async fn ai_agent_stream(
state: state.app.clone(),
connection_id: body.connection_id,
database: body.database,
selected_databases: body.selected_databases,
schema: body.schema,
db_type: parsed_db_type,
cli_mcp_server_command: None,
+2
View File
@@ -333,6 +333,7 @@ pub async fn ai_agent_stream(
confirmed_connection_id: Option<String>,
confirmed_database: Option<String>,
confirmed_schema: Option<String>,
selected_databases: Option<Vec<String>>,
) -> Result<String, String> {
let mut request = resolve_cli_provider_request(request);
merge_global_max_retries(
@@ -387,6 +388,7 @@ pub async fn ai_agent_stream(
state: state.inner().clone(),
connection_id,
database,
selected_databases: selected_databases.unwrap_or_default(),
schema,
db_type: parsed_db_type,
cli_mcp_server_command,