From 7a0f89d3ad793818c71b0e072f0c8fb7a96d41b5 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Wed, 10 Jun 2026 18:19:53 +0800
Subject: [PATCH] fix(jdbc): use schema tree for inferred db2

---
 apps/desktop/src/lib/jdbcDialect.ts    | 6 ++++--
 packages/app-tests/jdbcDialect.test.ts | 4 ++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/apps/desktop/src/lib/jdbcDialect.ts b/apps/desktop/src/lib/jdbcDialect.ts
index 5ffa39d7d3..0105635a7f 100644
--- a/apps/desktop/src/lib/jdbcDialect.ts
+++ b/apps/desktop/src/lib/jdbcDialect.ts
@@ -1,5 +1,5 @@
 import type { ConnectionConfig, DatabaseType } from "@/types/database";
-import { isSchemaAware, usesDatabaseObjectTreeMode } from "@/lib/databaseFeatureSupport";
+import { isSchemaAware, usesDatabaseObjectTreeMode, usesTreeSchemaMode } from "@/lib/databaseFeatureSupport";
 
 const JDBC_DIALECT_MATCHERS: Array<{ type: DatabaseType; patterns: RegExp[] }> = [
   { type: "databend", patterns: [/jdbc:databend:/i, /com\.databend\.jdbc\.DatabendDriver/i, /databend-jdbc/i] },
@@ -34,8 +34,10 @@ export function connectionUsesDatabaseObjectTreeMode(connection?: Pick<Connectio
   if (!connection) return false;
   if (connection.db_type !== "jdbc") return usesDatabaseObjectTreeMode(connection.db_type);
   const dialect = inferJdbcDialect(connection);
+  if (!dialect) return true;
   if (dialect === "hive" || dialect === "trino") return false;
-  return true;
+  if (dialect === "databend") return true;
+  return !usesTreeSchemaMode(dialect);
 }
 
 export function connectionUsesSchemaExecutionContext(connection?: Pick<ConnectionConfig, "db_type" | "connection_string" | "jdbc_driver_class" | "jdbc_driver_paths">): boolean {
diff --git a/packages/app-tests/jdbcDialect.test.ts b/packages/app-tests/jdbcDialect.test.ts
index ab02d1da0a..0ffc6b4caa 100644
--- a/packages/app-tests/jdbcDialect.test.ts
+++ b/packages/app-tests/jdbcDialect.test.ts
@@ -37,11 +37,15 @@ test("effective database type keeps non-JDBC types and enables compatible JDBC s
 test("JDBC tree shape follows the inferred driver dialect", () => {
   const kyuubi = { db_type: "jdbc" as const, jdbc_driver_class: "org.apache.kyuubi.jdbc.KyuubiHiveDriver" };
   const hive = { db_type: "jdbc" as const, jdbc_driver_class: "org.apache.hive.jdbc.HiveDriver" };
+  const db2 = { db_type: "jdbc" as const, connection_string: "jdbc:db2://db.example.com:50000/SAMPLE" };
 
   assert.equal(connectionUsesDatabaseObjectTreeMode(kyuubi), true);
   assert.equal(connectionObjectTreeQuerySchema(kyuubi, "test", undefined), "");
   assert.equal(connectionUsesDatabaseObjectTreeMode(hive), false);
   assert.equal(connectionObjectTreeQuerySchema(hive, "spark_catalog", "test"), "test");
+  assert.equal(connectionUsesDatabaseObjectTreeMode(db2), false);
+  assert.equal(connectionObjectTreeQuerySchema(db2, "SAMPLE", "APP"), "APP");
+  assert.equal(connectionObjectTreeNodeSchema(db2, "SAMPLE", "APP"), "APP");
   assert.equal(
     qualifiedTableName({
       databaseType: effectiveDatabaseTypeForConnection(hive),