diff --git a/crates/dbx-core/src/sql_dialect.rs b/crates/dbx-core/src/sql_dialect.rs index d6435788c1..ba5175bd01 100644 --- a/crates/dbx-core/src/sql_dialect.rs +++ b/crates/dbx-core/src/sql_dialect.rs @@ -727,6 +727,31 @@ mod tests { ); } + #[test] + fn builds_iris_table_data_sql_with_literal_top_and_quoted_object() { + let sql = build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Iris), + schema: Some("Ens".to_string()), + table_name: "AlarmResponse".to_string(), + primary_keys: vec!["ID".to_string()], + columns: vec!["ID".to_string(), "Status".to_string()], + fallback_order_columns: Vec::new(), + order_by: Some("\"Status\" DESC".to_string()), + limit: Some(25), + offset: None, + where_input: Some("WHERE \"Status\" = 'Open'".to_string()), + include_row_id: false, + }); + + assert_eq!( + sql, + "SELECT TOP 25 * FROM \"Ens\".\"AlarmResponse\" WHERE (\"Status\" = 'Open') ORDER BY \"Status\" DESC" + ); + assert!(!sql.contains("?")); + assert!(!sql.contains(":%qpar")); + assert!(!sql.contains(" LIMIT ")); + } + #[test] fn builds_table_data_special_column_queries() { assert_eq!( 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 ce0f50433e..665491691f 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -45,12 +45,40 @@ 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, false, false, false, false); - private static final JdbcDriverQuirks USE_CATALOG_QUIRKS = new JdbcDriverQuirks(false, false, false, true, false); - private static final JdbcDriverQuirks KINGBASE_QUIRKS = new JdbcDriverQuirks(false, false, false, false, true); - private static final JdbcDriverQuirks YASHAN_QUIRKS = new JdbcDriverQuirks(true, true, false, false, false); - private static final JdbcDriverQuirks IRIS_QUIRKS = new JdbcDriverQuirks(true, false, true, false, false); - private static final JdbcDriverQuirks ORACLE_QUIRKS = new JdbcDriverQuirks(false, true, false, false, false); + private static final JdbcDriverQuirks DEFAULT_QUIRKS = new JdbcDriverQuirks( + false, + false, + false, + false, + false, + StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS + ); + private static final JdbcDriverQuirks USE_CATALOG_QUIRKS = DEFAULT_QUIRKS.withUseCatalogFallbackSql(true); + private static final JdbcDriverQuirks KINGBASE_QUIRKS = DEFAULT_QUIRKS.withIgnoreCatalogForSchemaMetadata(true); + private static final JdbcDriverQuirks YASHAN_QUIRKS = new JdbcDriverQuirks( + true, + true, + false, + false, + false, + StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS + ); + private static final JdbcDriverQuirks IRIS_QUIRKS = new JdbcDriverQuirks( + true, + false, + true, + false, + false, + StatementMaxRowsMode.READ_LOOP_ONLY + ); + private static final JdbcDriverQuirks ORACLE_QUIRKS = new JdbcDriverQuirks( + false, + true, + false, + false, + false, + StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS + ); private static final List DRIVER_QUIRK_RULES = List.of( new JdbcDriverQuirkRule("jdbc:mysql:", USE_CATALOG_QUIRKS), new JdbcDriverQuirkRule("jdbc:mariadb:", USE_CATALOG_QUIRKS), @@ -72,8 +100,35 @@ public final class DbxJdbcPlugin { boolean useOracleMetadata, boolean caseInsensitiveSchemaMetadata, boolean useCatalogFallbackSql, - boolean ignoreCatalogForSchemaMetadata + boolean ignoreCatalogForSchemaMetadata, + StatementMaxRowsMode statementMaxRowsMode ) { + JdbcDriverQuirks withUseCatalogFallbackSql(boolean value) { + return new JdbcDriverQuirks( + skipExecutionContext, + useOracleMetadata, + caseInsensitiveSchemaMetadata, + value, + ignoreCatalogForSchemaMetadata, + statementMaxRowsMode + ); + } + + JdbcDriverQuirks withIgnoreCatalogForSchemaMetadata(boolean value) { + return new JdbcDriverQuirks( + skipExecutionContext, + useOracleMetadata, + caseInsensitiveSchemaMetadata, + useCatalogFallbackSql, + value, + statementMaxRowsMode + ); + } + } + + enum StatementMaxRowsMode { + APPLY_STATEMENT_MAX_ROWS, + READ_LOOP_ONLY } private record JdbcDriverQuirkRule(String urlPrefix, JdbcDriverQuirks quirks) { @@ -282,8 +337,9 @@ public final class DbxJdbcPlugin { long start = System.nanoTime(); Connection conn = openConnection(connection); applyExecutionContext(connection, conn, database, schema); + JdbcDriverQuirks quirks = driverQuirks(connection); try (Statement statement = conn.createStatement()) { - applyStatementOptions(statement, maxRows, fetchSize, timeoutSecs); + applyStatementOptions(statement, maxRows, fetchSize, timeoutSecs, quirks); boolean hasResultSet = statement.execute(trimStatementSql(sql)); ObjectNode result = MAPPER.createObjectNode(); ArrayNode columns = MAPPER.createArrayNode(); @@ -422,9 +478,17 @@ public final class DbxJdbcPlugin { return result; } - private static void applyStatementOptions(Statement statement, int maxRows, int fetchSize, int timeoutSecs) + private static void applyStatementOptions( + Statement statement, + int maxRows, + int fetchSize, + int timeoutSecs, + JdbcDriverQuirks quirks + ) throws SQLException { - statement.setMaxRows((int) Math.min(Integer.MAX_VALUE, (long) maxRows + 1L)); + if (quirks.statementMaxRowsMode() == StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS) { + statement.setMaxRows((int) Math.min(Integer.MAX_VALUE, (long) maxRows + 1L)); + } if (fetchSize > 0) { try { statement.setFetchSize(fetchSize); 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 64a24f495d..85804fb998 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -8,6 +8,10 @@ import org.junit.jupiter.api.Test; import java.sql.Connection; import java.sql.DriverManager; import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; import java.util.Properties; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -246,15 +250,73 @@ final class DbxJdbcPluginTest { assertEquals(true, DbxJdbcPlugin.driverQuirks(iris).skipExecutionContext()); assertEquals(false, DbxJdbcPlugin.driverQuirks(iris).useOracleMetadata()); assertEquals(true, DbxJdbcPlugin.driverQuirks(iris).caseInsensitiveSchemaMetadata()); + assertEquals( + DbxJdbcPlugin.StatementMaxRowsMode.READ_LOOP_ONLY, + DbxJdbcPlugin.driverQuirks(iris).statementMaxRowsMode() + ); assertEquals(false, DbxJdbcPlugin.driverQuirks(h2).skipExecutionContext()); assertEquals(false, DbxJdbcPlugin.driverQuirks(h2).useOracleMetadata()); assertEquals(false, DbxJdbcPlugin.driverQuirks(h2).caseInsensitiveSchemaMetadata()); assertEquals(false, DbxJdbcPlugin.driverQuirks(h2).useCatalogFallbackSql()); + assertEquals( + DbxJdbcPlugin.StatementMaxRowsMode.APPLY_STATEMENT_MAX_ROWS, + DbxJdbcPlugin.driverQuirks(h2).statementMaxRowsMode() + ); assertEquals(true, DbxJdbcPlugin.driverQuirks(mysql).useCatalogFallbackSql()); assertEquals(true, DbxJdbcPlugin.driverQuirks(kingbase).ignoreCatalogForSchemaMetadata()); assertEquals(true, DbxJdbcPlugin.driverQuirks(kyuubi).useCatalogFallbackSql()); } + @Test + void irisStatementOptionsSkipDriverMaxRowsRewrite() throws Exception { + Method method = DbxJdbcPlugin.class.getDeclaredMethod( + "applyStatementOptions", + Statement.class, + int.class, + int.class, + int.class, + DbxJdbcPlugin.JdbcDriverQuirks.class + ); + method.setAccessible(true); + JsonNode iris = MAPPER.readTree(""" + { + "connection_string": "jdbc:IRIS://127.0.0.1:1972/USER" + } + """); + List calls = new ArrayList<>(); + + method.invoke(null, recordingStatement(calls), 100, 50, 30, DbxJdbcPlugin.driverQuirks(iris)); + + assertFalse(calls.contains("setMaxRows"), calls.toString()); + assertEquals(true, calls.contains("setFetchSize")); + assertEquals(true, calls.contains("setQueryTimeout")); + } + + @Test + void defaultStatementOptionsApplyDriverMaxRowsProtection() throws Exception { + Method method = DbxJdbcPlugin.class.getDeclaredMethod( + "applyStatementOptions", + Statement.class, + int.class, + int.class, + int.class, + DbxJdbcPlugin.JdbcDriverQuirks.class + ); + method.setAccessible(true); + JsonNode h2 = MAPPER.readTree(""" + { + "connection_string": "jdbc:h2:mem:dbx_quirks" + } + """); + List calls = new ArrayList<>(); + + method.invoke(null, recordingStatement(calls), 100, 50, 30, DbxJdbcPlugin.driverQuirks(h2)); + + assertEquals(true, calls.contains("setMaxRows")); + assertEquals(true, calls.contains("setFetchSize")); + assertEquals(true, calls.contains("setQueryTimeout")); + } + @Test void schemaDisplayNamePrefersMixedCaseOverAllUppercaseDuplicate() { assertEquals(true, DbxJdbcPlugin.preferSchemaDisplayName("SQLUSER", "SQLUser")); @@ -488,6 +550,23 @@ final class DbxJdbcPluginTest { """.formatted(CONNECTION)); } + private static Statement recordingStatement(List calls) { + return (Statement) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Statement.class }, + (proxy, method, args) -> { + calls.add(method.getName()); + Class returnType = method.getReturnType(); + if (returnType == boolean.class) return false; + if (returnType == int.class) return 0; + if (returnType == long.class) return 0L; + if (returnType == float.class) return 0f; + if (returnType == double.class) return 0d; + return null; + } + ); + } + private static JsonNode request(String method, String params) throws Exception { Method handleLine = DbxJdbcPlugin.class.getDeclaredMethod("handleLine", String.class); handleLine.setAccessible(true);