github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 03:36:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(drivers): add JDBC quirks and RedisJSON previews

Browse Source
This commit is contained in:
t8y2
2026-05-18 14:07:14 +08:00
parent e1930550e2
commit 57dea3cb6d
3 changed files with 171 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/redis_driver.rs
+78 -4
View File
@@ -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<serde_json::Value, String> {
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::<String>();
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<serde_json::Value>
#[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"), "");
}
}
plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java
+51 -16
View File
@@ -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<JdbcDriverQuirkRule> 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);
plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java
+42
View File
@@ -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();