From 0655d309d88de97fcf2ba32b09b1a2657e6c3a2c Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Fri, 22 May 2026 14:52:12 +0800
Subject: [PATCH] fix(db2): use schema tree mode for sidebar object loading

---
 apps/desktop/src/lib/databaseCapabilitySets.ts  | 1 +
 crates/dbx-core/src/lib.rs                      | 1 +
 packages/app-tests/databaseCapabilities.test.ts | 5 +++++
 src-tauri/src/commands/mod.rs                   | 1 +
 src-tauri/src/lib.rs                            | 1 +
 5 files changed, 9 insertions(+)

diff --git a/apps/desktop/src/lib/databaseCapabilitySets.ts b/apps/desktop/src/lib/databaseCapabilitySets.ts
index 1ce433216..1f0c2a21d 100644
--- a/apps/desktop/src/lib/databaseCapabilitySets.ts
+++ b/apps/desktop/src/lib/databaseCapabilitySets.ts
@@ -183,6 +183,7 @@ export const TREE_SCHEMA_TYPES = new Set<DatabaseType>([
   "postgres",
   "redshift",
   "sqlserver",
+  "db2",
   "gaussdb",
   "kingbase",
   "highgo",
diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs
index 7ef109941..850e8b36d 100644
--- a/crates/dbx-core/src/lib.rs
+++ b/crates/dbx-core/src/lib.rs
@@ -23,6 +23,7 @@ pub mod table_import;
 pub mod transfer;
 pub mod types;
 pub mod update;
+pub mod xlsx_export;
 
 pub const R2_CDN_BASE: &str = "https://dl.dbxio.com/";
 
diff --git a/packages/app-tests/databaseCapabilities.test.ts b/packages/app-tests/databaseCapabilities.test.ts
index 207fdc9fc..e038d7873 100644
--- a/packages/app-tests/databaseCapabilities.test.ts
+++ b/packages/app-tests/databaseCapabilities.test.ts
@@ -24,6 +24,11 @@ test("treats Trino catalogs as schema tree roots", () => {
   assert.equal(TREE_SCHEMA_TYPES.has("trino"), true);
 });
 
+test("treats DB2 databases as schema tree roots", () => {
+  assert.equal(TREE_SCHEMA_TYPES.has("db2"), true);
+  assert.equal(usesTreeSchemaMode("db2"), true);
+});
+
 test("treats TDengine databases as schema tree roots and agent driver databases", () => {
   assert.equal(TREE_SCHEMA_TYPES.has("tdengine"), true);
   assert.equal(SCHEMA_AWARE_TYPES.has("tdengine"), true);
diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs
index 8d15dce38..267ca8925 100644
--- a/src-tauri/src/commands/mod.rs
+++ b/src-tauri/src/commands/mod.rs
@@ -22,3 +22,4 @@ pub mod system_fonts;
 pub mod table_import;
 pub mod transfer;
 pub mod update;
+pub mod xlsx_export;
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
index a9560e672..bb2ce1b25 100644
--- a/src-tauri/src/lib.rs
+++ b/src-tauri/src/lib.rs
@@ -337,6 +337,7 @@ pub fn run() {
             commands::transfer::cancel_transfer,
             commands::database_export::export_database_sql,
             commands::database_export::cancel_database_export,
+            commands::xlsx_export::export_query_result_xlsx,
             commands::agents::list_installed_agents,
             commands::agents::list_installed_agents_local,
             commands::agents::install_agent,