github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-20 00:17:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(iris): detect editable table primary keys

Browse Source
This commit is contained in:
t8y2
2026-08-01 00:06:33 +08:00
parent 290693e0fa
commit 1303466947
2 changed files with 161 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/drivers/iris/src/main/java/com/dbx/agent/iris/IrisAgent.java
+50 -9
View File
@@ -7,6 +7,7 @@ import com.dbx.agent.MultiSessionJsonRpcServer;
import com.dbx.agent.StandardJdbcMetadata;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
@@ -41,7 +42,6 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
static List<ColumnInfo> irisColumns(Connection conn, String schema, String table) {
return unchecked(() -> {
Set<String> primaryKeys = irisPrimaryKeys(conn, schema, table);
List<ColumnInfo> columns = new ArrayList<>();
try (PreparedStatement stmt = conn.prepareStatement(columnSql(schema))) {
int index = 1;
@@ -60,7 +60,7 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
stringOrDefault(rs, "DATA_TYPE", "Unknown"),
!"NO".equalsIgnoreCase(stringOrNull(rs, "IS_NULLABLE")),
stringOrNull(rs, "COLUMN_DEFAULT"),
primaryKeys.contains(name.toUpperCase(Locale.ROOT)),
"YES".equalsIgnoreCase(stringOrNull(rs, "PRIMARY_KEY")),
null,
null,
intOrNull(rs, "NUMERIC_PRECISION"),
@@ -70,6 +70,9 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
}
}
}
if (!hasPrimaryKey(columns)) {
markPrimaryKeys(columns, irisPrimaryKeys(conn, schema, table));
}
return columns;
});
}
@@ -87,6 +90,23 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
}
private static Set<String> irisPrimaryKeys(Connection conn, String schema, String table) {
Set<String> keys = jdbcPrimaryKeys(conn, schema, table);
return keys.isEmpty() ? nativePrimaryKeys(conn, schema, table) : keys;
}
private static Set<String> jdbcPrimaryKeys(Connection conn, String schema, String table) {
Set<String> keys = new LinkedHashSet<>();
try {
DatabaseMetaData meta = conn.getMetaData();
try (ResultSet rs = meta.getPrimaryKeys(null, isBlank(schema) ? null : schema.trim(), table)) {
addPrimaryKeys(keys, rs);
}
} catch (Exception ignored) {
}
return keys;
}
private static Set<String> nativePrimaryKeys(Connection conn, String schema, String table) {
Set<String> keys = new LinkedHashSet<>();
try (PreparedStatement stmt = conn.prepareStatement(primaryKeySql(schema))) {
int index = 1;
@@ -95,12 +115,7 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
}
stmt.setString(index, table);
try (ResultSet rs = stmt.executeQuery()) {
while (rs.next()) {
String column = rs.getString("COLUMN_NAME");
if (!isBlank(column)) {
keys.add(column.toUpperCase(Locale.ROOT));
}
}
addPrimaryKeys(keys, rs);
}
} catch (Exception ignored) {
// Primary key metadata is optional for display; column listing must keep working.
@@ -111,7 +126,7 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
private static String columnSql(String schema) {
StringBuilder sql = new StringBuilder(
"SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT, "
+ "NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH "
+ "PRIMARY_KEY, NUMERIC_PRECISION, NUMERIC_SCALE, CHARACTER_MAXIMUM_LENGTH "
+ "FROM INFORMATION_SCHEMA.COLUMNS WHERE "
);
if (!isBlank(schema)) {
@@ -121,6 +136,32 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
return sql.toString();
}
private static void addPrimaryKeys(Set<String> keys, ResultSet rs) throws Exception {
while (rs.next()) {
String column = rs.getString("COLUMN_NAME");
if (!isBlank(column)) {
keys.add(column.toUpperCase(Locale.ROOT));
}
}
}
private static boolean hasPrimaryKey(List<ColumnInfo> columns) {
for (ColumnInfo column : columns) {
if (column.getIs_primary_key()) {
return true;
}
}
return false;
}
private static void markPrimaryKeys(List<ColumnInfo> columns, Set<String> primaryKeys) {
for (ColumnInfo column : columns) {
if (primaryKeys.contains(column.getName().toUpperCase(Locale.ROOT))) {
column.setIs_primary_key(true);
}
}
}
private static String primaryKeySql(String schema) {
StringBuilder sql = new StringBuilder(
"SELECT k.COLUMN_NAME "
agents/drivers/iris/src/test/java/com/dbx/agent/iris/IrisAgentTest.java
+111 -11
View File
@@ -7,8 +7,10 @@ import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -44,9 +46,15 @@ class IrisAgentTest {
}
@Test
void readsColumnsWithNativeInformationSchemaQueries() {
List<String> sql = new ArrayList<>();
Connection conn = fakeConnection(sql);
void readsColumnsAndUsesInlinePrimaryKey() {
List<String> calls = new ArrayList<>();
Connection conn = fakeConnection(
calls,
true,
Collections.emptyList(),
null,
Collections.emptyList()
);
List<ColumnInfo> columns = IrisAgent.irisColumns(conn, "SQLUser", "People");
@@ -57,42 +65,134 @@ class IrisAgentTest {
assertTrue(columns.get(0).getIs_primary_key());
assertEquals("NAME", columns.get(1).getName());
assertEquals(Integer.valueOf(64), columns.get(1).getCharacter_maximum_length());
assertFalse(sql.get(0).contains("LIMIT"));
assertTrue(sql.get(0).contains("INFORMATION_SCHEMA.KEY_COLUMN_USAGE"));
assertTrue(sql.get(1).contains("INFORMATION_SCHEMA.COLUMNS"));
assertEquals(1, calls.size());
assertFalse(calls.get(0).contains("LIMIT"));
assertTrue(calls.get(0).contains("PRIMARY_KEY"));
assertTrue(calls.get(0).contains("INFORMATION_SCHEMA.COLUMNS"));
}
private static Connection fakeConnection(List<String> sqlLog) {
@Test
void usesJdbcPrimaryKeysWhenColumnRowsHaveNoPrimaryKey() {
List<String> calls = new ArrayList<>();
Connection conn = fakeConnection(
calls,
false,
List.of(Map.of("COLUMN_NAME", "ID")),
null,
List.of(Map.of("COLUMN_NAME", "NATIVE_ID"))
);
List<ColumnInfo> columns = IrisAgent.irisColumns(conn, "SQLUser", "People");
assertTrue(columns.get(0).getIs_primary_key());
assertEquals(2, calls.size());
assertTrue(calls.get(0).contains("INFORMATION_SCHEMA.COLUMNS"));
assertEquals("JDBC_PRIMARY_KEYS", calls.get(1));
}
@Test
void fallsBackToNativePrimaryKeysWhenJdbcMetadataFails() {
List<String> calls = new ArrayList<>();
Connection conn = fakeConnection(
calls,
false,
Collections.emptyList(),
new SQLException("IRIS metadata unavailable"),
List.of(Map.of("COLUMN_NAME", "ID"))
);
List<ColumnInfo> columns = IrisAgent.irisColumns(conn, "SQLUser", "People");
assertTrue(columns.get(0).getIs_primary_key());
assertTrue(calls.get(0).contains("INFORMATION_SCHEMA.COLUMNS"));
assertEquals("JDBC_PRIMARY_KEYS", calls.get(1));
assertTrue(calls.get(2).contains("INFORMATION_SCHEMA.KEY_COLUMN_USAGE"));
}
@Test
void leavesColumnsWithoutPrimaryKeysWhenBothSourcesAreEmpty() {
List<String> calls = new ArrayList<>();
Connection conn = fakeConnection(
calls,
false,
Collections.emptyList(),
null,
Collections.emptyList()
);
List<ColumnInfo> columns = IrisAgent.irisColumns(conn, "SQLUser", "People");
assertFalse(columns.get(0).getIs_primary_key());
assertFalse(columns.get(1).getIs_primary_key());
assertTrue(calls.get(2).contains("INFORMATION_SCHEMA.KEY_COLUMN_USAGE"));
}
private static Connection fakeConnection(
List<String> callLog,
boolean inlinePrimaryKey,
List<Map<String, Object>> jdbcPrimaryKeys,
SQLException jdbcPrimaryKeyFailure,
List<Map<String, Object>> nativePrimaryKeys
) {
return proxy(Connection.class, (method, args) -> {
if ("prepareStatement".equals(method.getName())) {
String sql = (String) args[0];
sqlLog.add(sql);
return fakeStatement(sql);
callLog.add(sql);
return fakeStatement(sql, inlinePrimaryKey, nativePrimaryKeys);
}
if ("getMetaData".equals(method.getName())) {
return fakeMetadata(callLog, jdbcPrimaryKeys, jdbcPrimaryKeyFailure);
}
return defaultValue(method.getReturnType());
});
}
private static DatabaseMetaData fakeMetadata(
List<String> callLog,
List<Map<String, Object>> jdbcPrimaryKeys,
SQLException jdbcPrimaryKeyFailure
) {
return proxy(DatabaseMetaData.class, (method, args) -> {
if ("getPrimaryKeys".equals(method.getName())) {
callLog.add("JDBC_PRIMARY_KEYS");
assertEquals(null, args[0]);
assertEquals("SQLUser", args[1]);
assertEquals("People", args[2]);
if (jdbcPrimaryKeyFailure != null) {
throw jdbcPrimaryKeyFailure;
}
return fakeResultSet(jdbcPrimaryKeys);
}
if ("getColumns".equals(method.getName())) {
throw new AssertionError("IRIS columns must not use JDBC DatabaseMetaData.getColumns()");
}
return defaultValue(method.getReturnType());
});
}
private static PreparedStatement fakeStatement(String sql) {
private static PreparedStatement fakeStatement(
String sql,
boolean inlinePrimaryKey,
List<Map<String, Object>> nativePrimaryKeys
) {
return proxy(PreparedStatement.class, (method, args) -> {
if ("executeQuery".equals(method.getName())) {
if (sql.contains("KEY_COLUMN_USAGE")) {
return fakeResultSet(List.of(Map.of("COLUMN_NAME", "ID")));
return fakeResultSet(nativePrimaryKeys);
}
return fakeResultSet(List.of(
Map.of(
"COLUMN_NAME", "ID",
"DATA_TYPE", "INTEGER",
"IS_NULLABLE", "NO",
"PRIMARY_KEY", inlinePrimaryKey ? "YES" : "NO",
"NUMERIC_PRECISION", 10
),
Map.of(
"COLUMN_NAME", "NAME",
"DATA_TYPE", "VARCHAR",
"IS_NULLABLE", "YES",
"PRIMARY_KEY", "NO",
"CHARACTER_MAXIMUM_LENGTH", 64
)
));