From 57dea3cb6d34ef3cbc93598c77844d4d5e92ce30 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 18 May 2026 14:07:14 +0800 Subject: [PATCH] feat(drivers): add JDBC quirks and RedisJSON previews --- crates/dbx-core/src/db/redis_driver.rs | 82 ++++++++++++++++++- .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 67 +++++++++++---- .../java/app/dbx/jdbc/DbxJdbcPluginTest.java | 42 ++++++++++ 3 files changed, 171 insertions(+), 20 deletions(-) diff --git a/crates/dbx-core/src/db/redis_driver.rs b/crates/dbx-core/src/db/redis_driver.rs index 11016534ef..1640f7efad 100644 --- a/crates/dbx-core/src/db/redis_driver.rs +++ b/crates/dbx-core/src/db/redis_driver.rs @@ -254,6 +254,39 @@ pub fn redis_command_raw_to_json(value: RedisRawValue) -> serde_json::Value { } } +pub fn is_redis_json_type(key_type: &str) -> bool { + matches!(key_type.to_ascii_uppercase().as_str(), "REJSON-RL" | "JSON") +} + +pub fn redis_json_raw_to_json(value: RedisRawValue) -> Result { + match redis_raw_to_json(value) { + serde_json::Value::Null => Ok(serde_json::Value::Null), + serde_json::Value::String(text) => { + serde_json::from_str(&text).map_err(|e| format!("Invalid RedisJSON value: {e}")) + } + other => Ok(other), + } +} + +pub fn redis_json_value_preview(value: &serde_json::Value) -> String { + const MAX_PREVIEW_LEN: usize = 160; + let text = serde_json::to_string(value).unwrap_or_else(|_| value.to_string()); + if text.chars().count() <= MAX_PREVIEW_LEN { + return text; + } + let mut preview = text.chars().take(MAX_PREVIEW_LEN).collect::(); + preview.push('…'); + preview +} + +pub fn redis_key_value_preview(key_type: &str) -> String { + if is_redis_json_type(key_type) { + "{...}".to_string() + } else { + String::new() + } +} + pub async fn flush_db(con: &mut redis::aio::MultiplexedConnection) -> Result<(), String> { redis::cmd("FLUSHDB").query_async::<()>(con).await.map_err(|e| e.to_string()) } @@ -308,13 +341,14 @@ pub async fn scan_keys_page( let mut result = Vec::with_capacity(keys.len()); for (index, key) in keys.iter().enumerate() { + let key_type = key_types.get(index).cloned().unwrap_or_else(|| "unknown".to_string()); result.push(RedisKeyInfo { key_display: redis_key_bytes_to_display(key), key_raw: redis_key_bytes_to_raw(key), - key_type: key_types.get(index).cloned().unwrap_or_else(|| "unknown".to_string()), + key_type, ttl: -2, size: 0, - value_preview: String::new(), + value_preview: redis_key_value_preview(key_types.get(index).map(String::as_str).unwrap_or("unknown")), }); } Ok(RedisScanResult { cursor: next_cursor, keys: result, total_keys }) @@ -358,6 +392,11 @@ pub async fn get_value(con: &mut redis::aio::MultiplexedConnection, key: &[u8]) (serde_json::Value::Array(items), false, Some(len), cursor) } "stream" => (get_stream_entries(con, key).await?, false, None, None), + key_type if is_redis_json_type(key_type) => { + let raw: RedisRawValue = + redis::cmd("JSON.GET").arg(key).query_async(con).await.map_err(|e| e.to_string())?; + (redis_json_raw_to_json(raw)?, false, None, None) + } _ => (serde_json::Value::Null, false, None, None), }; @@ -777,8 +816,9 @@ fn parse_scan_members(raw: RedisRawValue) -> Result<(u64, Vec #[cfg(test)] mod tests { use super::{ - classify_command, parse_command_argv, parse_database_count, parse_scan_keys, parse_stream_entries, - redis_command_raw_to_json, redis_key_bytes_to_display, redis_key_bytes_to_raw, redis_key_raw_to_bytes, + classify_command, is_redis_json_type, parse_command_argv, parse_database_count, parse_scan_keys, + parse_stream_entries, redis_command_raw_to_json, redis_json_raw_to_json, redis_json_value_preview, + redis_key_bytes_to_display, redis_key_bytes_to_raw, redis_key_raw_to_bytes, redis_key_value_preview, redis_raw_to_json, redis_value_contains_binary, RedisCommandSafety, RedisRawValue, }; @@ -930,4 +970,38 @@ mod tests { assert_eq!(redis_command_raw_to_json(raw), serde_json::json!(["OK", 2, null])); } + + #[test] + fn recognizes_redis_json_module_key_types() { + assert!(is_redis_json_type("ReJSON-RL")); + assert!(is_redis_json_type("json")); + assert!(!is_redis_json_type("string")); + } + + #[test] + fn parses_redis_json_get_bulk_string() { + let raw = bulk(r#"{"id":1,"embedding":[0.1,0.2],"meta":{"source":"test"}}"#); + + assert_eq!( + redis_json_raw_to_json(raw).unwrap(), + serde_json::json!({ + "id": 1, + "embedding": [0.1, 0.2], + "meta": { "source": "test" } + }) + ); + } + + #[test] + fn builds_compact_redis_json_value_preview() { + let value = serde_json::json!({ "id": 1, "embedding": [0.1, 0.2] }); + + assert_eq!(redis_json_value_preview(&value), r#"{"id":1,"embedding":[0.1,0.2]}"#); + } + + #[test] + fn uses_lightweight_redis_json_placeholder_for_key_scan_preview() { + assert_eq!(redis_key_value_preview("ReJSON-RL"), "{...}"); + assert_eq!(redis_key_value_preview("string"), ""); + } } diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 04c120e767..8627f4af76 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -38,10 +38,21 @@ import java.util.logging.Logger; public final class DbxJdbcPlugin { private static final ObjectMapper MAPPER = new ObjectMapper(); private static final int MAX_ROWS = 10_000; + private static final JdbcDriverQuirks DEFAULT_QUIRKS = new JdbcDriverQuirks(false); + private static final JdbcDriverQuirks YASHAN_QUIRKS = new JdbcDriverQuirks(true); + private static final List DRIVER_QUIRK_RULES = List.of( + new JdbcDriverQuirkRule("jdbc:yasdb:", YASHAN_QUIRKS) + ); private static String registeredDriverKey = ""; private static String sharedConnectionKey = ""; private static Connection sharedConnection; + record JdbcDriverQuirks(boolean skipExecutionContext) { + } + + private record JdbcDriverQuirkRule(String urlPrefix, JdbcDriverQuirks quirks) { + } + private DbxJdbcPlugin() { } @@ -192,7 +203,7 @@ public final class DbxJdbcPlugin { private static JsonNode executeQuery(JsonNode connection, String sql, String database, String schema) throws SQLException { long start = System.nanoTime(); Connection conn = openConnection(connection); - applyExecutionContext(conn, database, schema); + applyExecutionContext(connection, conn, database, schema); try (Statement statement = conn.createStatement()) { statement.setMaxRows(MAX_ROWS + 1); boolean hasResultSet = statement.execute(sql); @@ -232,7 +243,10 @@ public final class DbxJdbcPlugin { } } - private static void applyExecutionContext(Connection conn, String database, String schema) throws SQLException { + private static void applyExecutionContext(JsonNode connection, Connection conn, String database, String schema) throws SQLException { + if (driverQuirks(connection).skipExecutionContext()) { + return; + } if (database != null) { try { conn.setCatalog(database); @@ -247,27 +261,48 @@ public final class DbxJdbcPlugin { } } + static JdbcDriverQuirks driverQuirks(JsonNode connection) { + String url = optionalText(connection, "connection_string"); + for (JdbcDriverQuirkRule rule : DRIVER_QUIRK_RULES) { + if (urlMatchesPrefix(url, rule.urlPrefix())) { + return rule.quirks(); + } + } + return DEFAULT_QUIRKS; + } + + private static boolean urlMatchesPrefix(String url, String prefix) { + return url != null && url.regionMatches(true, 0, prefix, 0, prefix.length()); + } + private static JsonNode listDatabases(JsonNode connection) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); - try (ResultSet rs = conn.getMetaData().getCatalogs()) { - while (rs.next()) { - String name = rs.getString("TABLE_CAT"); - if (name != null && !name.isBlank()) { - ObjectNode item = MAPPER.createObjectNode(); - item.put("name", name); - result.add(item); - } - } - } - if (result.isEmpty() && conn.getCatalog() != null) { - ObjectNode item = MAPPER.createObjectNode(); - item.put("name", conn.getCatalog()); - result.add(item); + try (ResultSet rs = conn.getMetaData().getCatalogs()) { + while (rs.next()) { + String name = rs.getString("TABLE_CAT"); + addDatabase(result, name); } + } + addDatabase(result, optionalText(connection, "database")); + addDatabase(result, conn.getCatalog()); return result; } + private static void addDatabase(ArrayNode result, String name) { + if (name == null || name.isBlank()) { + return; + } + for (JsonNode item : result) { + if (name.equals(item.path("name").asText())) { + return; + } + } + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", name); + result.add(item); + } + private static JsonNode listSchemas(JsonNode connection, String database) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index c6cd7d0884..88724aea7d 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -47,6 +47,23 @@ final class DbxJdbcPluginTest { assertEquals("APP", response.path("result").path("rows").path(0).path(0).asText()); } + @Test + void driverQuirksDetectYashanJdbcUrl() throws Exception { + JsonNode yashan = MAPPER.readTree(""" + { + "connection_string": "jdbc:yasdb://172.26.128.159:20027/yasdb" + } + """); + JsonNode h2 = MAPPER.readTree(""" + { + "connection_string": "jdbc:h2:mem:dbx_quirks" + } + """); + + assertEquals(true, DbxJdbcPlugin.driverQuirks(yashan).skipExecutionContext()); + assertEquals(false, DbxJdbcPlugin.driverQuirks(h2).skipExecutionContext()); + } + @Test void listTablesFallsBackWhenCatalogFiltersEverything() throws Exception { request("executeQuery", """ @@ -74,6 +91,31 @@ final class DbxJdbcPluginTest { assertEquals("PEOPLE", response.path("result").path(0).path("name").asText()); } + @Test + void listDatabasesIncludesConfiguredDatabaseWhenDriverDoesNotReturnIt() throws Exception { + String connection = """ + { + "connection_string": "jdbc:h2:mem:dbx_catalog;DB_CLOSE_DELAY=-1", + "username": "sa", + "database": "DBX_DEMO" + } + """; + + JsonNode response = request("listDatabases", """ + { "connection": %s } + """.formatted(connection)); + + assertFalse(response.has("error"), response.toString()); + boolean found = false; + for (JsonNode database : response.path("result")) { + if ("DBX_DEMO".equals(database.path("name").asText())) { + found = true; + break; + } + } + assertEquals(true, found); + } + @Test void listObjectsAcceptsCamelCaseMethodAndFallsBackWhenCatalogFiltersEverything() throws Exception { createPeopleTable();