diff --git a/crates/dbx-core/src/sql_dialect/capabilities.rs b/crates/dbx-core/src/sql_dialect/capabilities.rs index 6e2112cf6a..6913ff2ee7 100644 --- a/crates/dbx-core/src/sql_dialect/capabilities.rs +++ b/crates/dbx-core/src/sql_dialect/capabilities.rs @@ -65,12 +65,13 @@ pub fn uses_fetch_first(database_type: DatabaseType) -> bool { pub fn pagination_strategy(database_type: Option, context: PaginationContext) -> TablePaginationStrategy { match database_type { Some(DatabaseType::Jdbc) => TablePaginationStrategy::AgentMaxRows, - Some(DatabaseType::Oracle) - if matches!(context, PaginationContext::TablePreview | PaginationContext::UserQuery) => - { - TablePaginationStrategy::Unbounded + Some(DatabaseType::Oracle) if matches!(context, PaginationContext::TablePreview) => { + TablePaginationStrategy::Rownum } - Some(DatabaseType::Oracle) => TablePaginationStrategy::FetchFirst, + Some(DatabaseType::Oracle) if matches!(context, PaginationContext::BoundedRead) => { + TablePaginationStrategy::FetchFirst + } + Some(DatabaseType::Oracle) => TablePaginationStrategy::Unbounded, Some(DatabaseType::Dameng) => TablePaginationStrategy::FetchFirst, Some(DatabaseType::Db2) => TablePaginationStrategy::Db2FetchFirst, Some(DatabaseType::SqlServer) => TablePaginationStrategy::SqlServerTop, diff --git a/crates/dbx-core/src/sql_dialect/table_select.rs b/crates/dbx-core/src/sql_dialect/table_select.rs index a25bb730f7..46574e9143 100644 --- a/crates/dbx-core/src/sql_dialect/table_select.rs +++ b/crates/dbx-core/src/sql_dialect/table_select.rs @@ -35,6 +35,16 @@ pub fn build_table_data_select_sql(options: TableDataSelectSqlOptions) -> String } else { build_select_columns(database_type, &options.columns) }; + let rownum_select_columns = quoted_table_columns_or_star(database_type, &options.columns); + let page_select_columns = if options.include_row_id && database_type == Some(DatabaseType::Oracle) { + if options.columns.is_empty() { + "*".to_string() + } else { + format!("\"{DBX_ROWID_COLUMN}\", {rownum_select_columns}") + } + } else { + rownum_select_columns.clone() + }; let table_alias = if options.include_row_id && database_type.is_some_and(uses_fetch_first) { format!("{table} t") } else { @@ -70,7 +80,20 @@ pub fn build_table_data_select_sql(options: TableDataSelectSqlOptions) -> String ) } TablePaginationStrategy::Rownum => { - build_rownum_table_select_sql(&table_alias, &where_clause, &order, &select_columns, limit) + let rownum_inner_select_columns = if options.include_row_id && database_type == Some(DatabaseType::Oracle) { + &select_columns + } else { + &rownum_select_columns + }; + build_rownum_table_select_sql( + &table_alias, + &where_clause, + &order, + rownum_inner_select_columns, + &page_select_columns, + limit, + options.offset.unwrap_or(0), + ) } TablePaginationStrategy::Unbounded => { format!("SELECT {select_columns} FROM {table_alias}{where_clause}{order}") @@ -108,16 +131,7 @@ pub fn build_table_data_select_sql(options: TableDataSelectSqlOptions) -> String pub fn build_table_select_sql(options: TableSelectSqlOptions<'_>) -> String { let database_type = options.database_type; let table = qualified_table_name(database_type, options.schema, options.table_name); - let select_columns = if options.columns.is_empty() { - "*".to_string() - } else { - options - .columns - .iter() - .map(|column| quote_table_identifier(database_type, column)) - .collect::>() - .join(", ") - }; + let select_columns = quoted_table_columns_or_star(database_type, options.columns); let order_by = if options.order_columns.is_empty() { String::new() } else { @@ -138,7 +152,9 @@ pub fn build_table_select_sql(options: TableSelectSqlOptions<'_>) -> String { TablePaginationStrategy::InformixFirst => { format!("SELECT FIRST {limit} {select_columns} FROM {table}{order_by}") } - TablePaginationStrategy::Rownum => build_rownum_table_select_sql(&table, "", &order_by, &select_columns, limit), + TablePaginationStrategy::Rownum => { + build_rownum_table_select_sql(&table, "", &order_by, &select_columns, &select_columns, limit, 0) + } TablePaginationStrategy::Db2FetchFirst | TablePaginationStrategy::FetchFirst => { format!("SELECT {select_columns} FROM {table}{order_by} FETCH FIRST {limit} ROWS ONLY") } @@ -161,15 +177,32 @@ fn informix_row_limit_clause(limit: usize, offset: usize) -> String { } } +fn quoted_table_columns_or_star(database_type: Option, columns: &[String]) -> String { + if columns.is_empty() { + return "*".to_string(); + } + columns.iter().map(|column| quote_table_identifier(database_type, column)).collect::>().join(", ") +} + fn build_rownum_table_select_sql( table: &str, where_clause: &str, order: &str, - select_columns: &str, + inner_select_columns: &str, + outer_select_columns: &str, limit: usize, + offset: usize, ) -> String { - let inner_select = format!("SELECT {select_columns} FROM {table}{where_clause}{order}"); - format!("SELECT {select_columns} FROM ({inner_select}) WHERE ROWNUM <= {limit}") + let inner_select = format!("SELECT {inner_select_columns} FROM {table}{where_clause}{order}"); + if offset == 0 { + return format!("SELECT {outer_select_columns} FROM ({inner_select}) WHERE ROWNUM <= {limit}"); + } + + let row_number_alias = quote_table_identifier(Some(DatabaseType::Oracle), "__dbx_row_num"); + let end = offset + limit; + format!( + "SELECT {outer_select_columns} FROM (SELECT dbx_inner.*, ROWNUM AS {row_number_alias} FROM ({inner_select}) dbx_inner WHERE ROWNUM <= {end}) WHERE {row_number_alias} > {offset}" + ) } pub(super) fn is_tdengine_tbname(database_type: Option, name: &str) -> bool { diff --git a/crates/dbx-core/src/sql_dialect/tests.rs b/crates/dbx-core/src/sql_dialect/tests.rs index 48bd49fa5f..697db110e1 100644 --- a/crates/dbx-core/src/sql_dialect/tests.rs +++ b/crates/dbx-core/src/sql_dialect/tests.rs @@ -63,7 +63,7 @@ fn maps_table_pagination_strategy_by_database_type() { assert_eq!(table_pagination_strategy(Some(DatabaseType::Informix)), TablePaginationStrategy::InformixFirst); assert_eq!(table_pagination_strategy(Some(DatabaseType::OceanbaseOracle)), TablePaginationStrategy::Rownum); assert_eq!(table_pagination_strategy(Some(DatabaseType::Questdb)), TablePaginationStrategy::QuestDbLimit); - assert_eq!(table_pagination_strategy(Some(DatabaseType::Oracle)), TablePaginationStrategy::Unbounded); + assert_eq!(table_pagination_strategy(Some(DatabaseType::Oracle)), TablePaginationStrategy::Rownum); assert_eq!( pagination_strategy(Some(DatabaseType::Oracle), PaginationContext::BoundedRead), TablePaginationStrategy::FetchFirst @@ -114,6 +114,17 @@ fn builds_select_sql_with_limit_syntax_for_database_type() { }), "SELECT \"id\", \"name\" FROM \"DB2INST1\".\"USERS\" ORDER BY \"id\" ASC FETCH FIRST 100 ROWS ONLY" ); + assert_eq!( + build_table_select_sql(TableSelectSqlOptions { + database_type: Some(DatabaseType::Oracle), + schema: Some("DBXTEST"), + table_name: "USERS", + columns: &columns, + order_columns: &keys, + limit: 100, + }), + "SELECT \"id\", \"name\" FROM (SELECT \"id\", \"name\" FROM \"DBXTEST\".\"USERS\" ORDER BY \"id\" ASC) WHERE ROWNUM <= 100" + ); assert_eq!( build_table_select_sql(TableSelectSqlOptions { database_type: Some(DatabaseType::OceanbaseOracle), @@ -304,6 +315,22 @@ fn builds_table_data_where_and_schema_queries() { }), "SELECT * FROM \"DB2INST1\".\"ORDERS\" WHERE (amount > 10) FETCH FIRST 50 ROWS ONLY" ); + assert_eq!( + build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Oracle), + schema: Some("DBXTEST".to_string()), + table_name: "ORDERS".to_string(), + primary_keys: vec!["ID".to_string()], + columns: vec!["ID".to_string(), "AMOUNT".to_string()], + fallback_order_columns: Vec::new(), + order_by: Some("\"ID\" ASC".to_string()), + limit: Some(50), + offset: Some(100), + where_input: Some("WHERE amount > 10".to_string()), + include_row_id: false, + }), + "SELECT \"ID\", \"AMOUNT\" FROM (SELECT dbx_inner.*, ROWNUM AS \"__dbx_row_num\" FROM (SELECT \"ID\", \"AMOUNT\" FROM \"DBXTEST\".\"ORDERS\" WHERE (amount > 10) ORDER BY \"ID\" ASC) dbx_inner WHERE ROWNUM <= 150) WHERE \"__dbx_row_num\" > 100" + ); assert_eq!( build_table_data_select_sql(TableDataSelectSqlOptions { database_type: Some(DatabaseType::OceanbaseOracle), @@ -557,7 +584,23 @@ fn builds_oracle_and_neo4j_table_data_queries() { where_input: None, include_row_id: true, }), - "SELECT ROWIDTOCHAR(t.ROWID) AS \"__DBX_ROWID\", t.* FROM \"DBXTEST\".\"DBX_LOAD_TABLE_006\" t" + "SELECT * FROM (SELECT ROWIDTOCHAR(t.ROWID) AS \"__DBX_ROWID\", t.* FROM \"DBXTEST\".\"DBX_LOAD_TABLE_006\" t) WHERE ROWNUM <= 100" + ); + assert_eq!( + build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Oracle), + schema: Some("DBXTEST".to_string()), + table_name: "DBX_LOAD_TABLE_006".to_string(), + primary_keys: vec![DBX_ROWID_COLUMN.to_string()], + columns: vec!["ID".to_string(), "NAME".to_string()], + fallback_order_columns: Vec::new(), + order_by: None, + limit: Some(100), + offset: None, + where_input: None, + include_row_id: true, + }), + "SELECT \"__DBX_ROWID\", \"ID\", \"NAME\" FROM (SELECT ROWIDTOCHAR(t.ROWID) AS \"__DBX_ROWID\", t.* FROM \"DBXTEST\".\"DBX_LOAD_TABLE_006\" t) WHERE ROWNUM <= 100" ); assert_eq!( build_table_data_select_sql(TableDataSelectSqlOptions {