From f659a74a87ef37f107cd2a0798d045542348ee30 Mon Sep 17 00:00:00 2001
From: LRcoding <38545827+LwClick@users.noreply.github.com>
Date: Tue, 7 Jul 2026 14:41:52 +0800
Subject: [PATCH] fix(jdbc): recover primary keys via catalog fallback

---
 .../com/dbx/agent/StandardJdbcMetadata.java   |  21 +++-
 .../dbx/agent/StandardJdbcMetadataTest.java   | 111 ++++++++++++++++++
 2 files changed, 131 insertions(+), 1 deletion(-)

diff --git a/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java b/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java
index 386e1fd6c7..fe395be688 100644
--- a/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java
+++ b/agents/common/src/main/java/com/dbx/agent/StandardJdbcMetadata.java
@@ -168,7 +168,11 @@ public final class StandardJdbcMetadata {
             Set<String> primaryKeys = primaryKeys(meta, null, schema, table);
             List<ColumnInfo> result = new ArrayList<>();
             appendColumns(result, meta, null, schema, table, primaryKeys);
-            if (result.isEmpty() && profile.getCatalogFallbackEnabled() && !configuredDatabase.trim().isEmpty()) {
+            if (!result.isEmpty() && primaryKeys.isEmpty() && profile.getCatalogFallbackEnabled() && hasConfiguredDatabase(configuredDatabase)) {
+                Set<String> fallbackPrimaryKeys = primaryKeys(meta, configuredDatabase, schema, table);
+                markPrimaryKeys(result, fallbackPrimaryKeys);
+            }
+            if (result.isEmpty() && profile.getCatalogFallbackEnabled() && hasConfiguredDatabase(configuredDatabase)) {
                 Set<String> fallbackPrimaryKeys = primaryKeys(meta, configuredDatabase, schema, table);
                 appendColumns(result, meta, configuredDatabase, schema, table, fallbackPrimaryKeys);
             }
@@ -522,6 +526,21 @@ public final class StandardJdbcMetadata {
         return keys;
     }
 
+    private static boolean hasConfiguredDatabase(String configuredDatabase) {
+        return configuredDatabase != null && !configuredDatabase.trim().isEmpty();
+    }
+
+    private static void markPrimaryKeys(List<ColumnInfo> columns, Set<String> primaryKeys) {
+        if (primaryKeys.isEmpty()) {
+            return;
+        }
+        for (ColumnInfo column : columns) {
+            if (primaryKeys.contains(column.getName())) {
+                column.setIs_primary_key(true);
+            }
+        }
+    }
+
     private static void appendSchemas(Set<String> names, ResultSet rs) throws Exception {
         try (ResultSet closeable = rs) {
             while (closeable.next()) {
diff --git a/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java b/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java
index 40bbbf95d0..34231659c8 100644
--- a/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java
+++ b/agents/common/src/test/java/com/dbx/agent/StandardJdbcMetadataTest.java
@@ -230,6 +230,64 @@ class StandardJdbcMetadataTest {
         assertEquals(Integer.valueOf(64), columns.get(1).getCharacter_maximum_length());
     }
 
+    @Test
+    void marksPrimaryKeysFromConfiguredCatalogWhenColumnsAlreadyLoaded() {
+        Connection conn = catalogPrimaryKeyFallbackConnection(
+            "TENANTDB",
+            rows(),
+            rows(row("COLUMN_NAME", "ID")),
+            orderColumns()
+        );
+
+        List<ColumnInfo> columns = StandardJdbcMetadata.INSTANCE.getColumns(conn, profile, "TENANTDB", "APP", "ORDERS");
+
+        assertEquals("ID", columns.get(0).getName());
+        assertTrue(columns.get(0).getIs_primary_key());
+        assertFalse(columns.get(0).getIs_nullable());
+        assertEquals("NAME", columns.get(1).getName());
+        assertFalse(columns.get(1).getIs_primary_key());
+        assertEquals(Integer.valueOf(64), columns.get(1).getCharacter_maximum_length());
+    }
+
+    @Test
+    void doesNotMarkPrimaryKeysFromConfiguredCatalogWhenFallbackUnavailable() {
+        JdbcAgentProfile fallbackDisabledProfile = new JdbcAgentProfile(
+            "example.Driver",
+            "jdbc:example://{host}:{port}/{database}",
+            0,
+            false,
+            Collections.emptySet(),
+            Collections.singletonList("TABLE"),
+            "\"",
+            "SET SCHEMA",
+            false,
+            false,
+            false,
+            false
+        );
+        Connection disabledConn = catalogPrimaryKeyFallbackConnection(
+            "TENANTDB",
+            rows(),
+            rows(row("COLUMN_NAME", "ID")),
+            orderColumns()
+        );
+
+        List<ColumnInfo> disabledColumns = StandardJdbcMetadata.INSTANCE.getColumns(disabledConn, fallbackDisabledProfile, "TENANTDB", "APP", "ORDERS");
+
+        assertFalse(disabledColumns.get(0).getIs_primary_key());
+
+        Connection emptyCatalogConn = catalogPrimaryKeyFallbackConnection(
+            "TENANTDB",
+            rows(),
+            rows(row("COLUMN_NAME", "ID")),
+            orderColumns()
+        );
+
+        List<ColumnInfo> emptyCatalogColumns = StandardJdbcMetadata.INSTANCE.getColumns(emptyCatalogConn, profile, " ", "APP", "ORDERS");
+
+        assertFalse(emptyCatalogColumns.get(0).getIs_primary_key());
+    }
+
     @Test
     void groupsIndexColumnsInOrdinalOrder() {
         Connection conn = connection(
@@ -490,6 +548,59 @@ class StandardJdbcMetadataTest {
         });
     }
 
+    private static Connection catalogPrimaryKeyFallbackConnection(
+        String fallbackCatalog,
+        ResultSet defaultPrimaryKeys,
+        ResultSet fallbackPrimaryKeys,
+        ResultSet columns
+    ) {
+        DatabaseMetaData meta = proxy(DatabaseMetaData.class, new MethodHandler() {
+            @Override
+            public Object handle(Method method, Object[] args) {
+                String name = method.getName();
+                if ("getPrimaryKeys".equals(name)) {
+                    return fallbackCatalog.equals(args[0]) ? fallbackPrimaryKeys : defaultPrimaryKeys;
+                }
+                if ("getColumns".equals(name)) {
+                    return columns;
+                }
+                if ("getCatalogs".equals(name)) {
+                    return rows();
+                }
+                return defaultValue(method.getReturnType());
+            }
+        });
+        return proxy(Connection.class, new MethodHandler() {
+            @Override
+            public Object handle(Method method, Object[] args) {
+                if ("getMetaData".equals(method.getName())) {
+                    return meta;
+                }
+                return defaultValue(method.getReturnType());
+            }
+        });
+    }
+
+    private static ResultSet orderColumns() {
+        return rows(row(
+            "COLUMN_NAME", "ID",
+            "TYPE_NAME", "INTEGER",
+            "NULLABLE", DatabaseMetaData.columnNoNulls,
+            "COLUMN_DEF", null,
+            "REMARKS", "identifier",
+            "COLUMN_SIZE", 10,
+            "DECIMAL_DIGITS", 0
+        ), row(
+            "COLUMN_NAME", "NAME",
+            "TYPE_NAME", "VARCHAR",
+            "NULLABLE", DatabaseMetaData.columnNullable,
+            "COLUMN_DEF", null,
+            "REMARKS", null,
+            "COLUMN_SIZE", 64,
+            "DECIMAL_DIGITS", 0
+        ));
+    }
+
     private static ResultSet rows(Map<String, Object>... rows) {
         return proxy(ResultSet.class, new MethodHandler() {
             private int index = -1;