diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 503089b46..b4b0a7f9c 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -2945,6 +2945,77 @@ pub async fn get_columns(pool: &Pool, schema: &str, table: &str) -> Result String { + format!("'{}'", value.replace('\'', "''")) +} + +fn redshift_columns_sql(schema: &str, table: &str) -> String { + format!( + "SELECT c.column_name, \ + CASE WHEN c.data_type = 'USER-DEFINED' THEN c.udt_name ELSE c.data_type END AS full_type, \ + c.is_nullable, \ + c.column_default, \ + CAST(c.numeric_precision AS varchar) AS numeric_precision, \ + CAST(c.numeric_scale AS varchar) AS numeric_scale, \ + CAST(c.character_maximum_length AS varchar) AS character_maximum_length \ + FROM information_schema.columns c \ + WHERE c.table_schema = {} AND c.table_name = {} \ + ORDER BY c.ordinal_position", + pg_quote_literal(schema), + pg_quote_literal(table) + ) +} + +fn query_result_text(row: &[serde_json::Value], index: usize) -> Option { + row.get(index).and_then(|value| match value { + serde_json::Value::String(value) => Some(value.clone()), + serde_json::Value::Number(value) => Some(value.to_string()), + serde_json::Value::Bool(value) => Some(value.to_string()), + serde_json::Value::Null | serde_json::Value::Array(_) | serde_json::Value::Object(_) => None, + }) +} + +fn query_result_i32(row: &[serde_json::Value], index: usize) -> Option { + query_result_text(row, index)?.parse().ok() +} + +fn redshift_columns_from_query_result(result: QueryResult) -> Vec { + result + .rows + .into_iter() + .filter_map(|row| { + Some(ColumnInfo { + name: query_result_text(&row, 0)?, + data_type: query_result_text(&row, 1).unwrap_or_default(), + is_nullable: query_result_text(&row, 2).is_none_or(|value| value.eq_ignore_ascii_case("YES")), + column_default: query_result_text(&row, 3), + is_primary_key: false, + extra: None, + comment: None, + numeric_precision: query_result_i32(&row, 4), + numeric_scale: query_result_i32(&row, 5), + character_maximum_length: query_result_i32(&row, 6), + enum_values: None, + ..Default::default() + }) + }) + .collect() +} + +pub async fn get_redshift_columns(pool: &Pool, schema: &str, table: &str) -> Result, String> { + let schema = if schema.is_empty() { "public" } else { schema }; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let result = execute_select_text( + &client, + &redshift_columns_sql(schema, table), + Instant::now(), + crate::query::MAX_ROWS, + None, + ) + .await?; + Ok(redshift_columns_from_query_result(result)) +} + pub(crate) fn pg_quote_ident(ident: &str) -> String { format!("\"{}\"", ident.replace('"', "\"\"")) } @@ -3016,6 +3087,7 @@ pub async fn execute_query_with_max_rows_and_cancel( cancel_token: Option, budget: DbOperationBudget, cancel_context: Option, + prefer_text_protocol: bool, ) -> Result { let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; let pg_cancel_token = client.cancel_token(); @@ -3025,7 +3097,7 @@ pub async fn execute_query_with_max_rows_and_cancel( cancel_token, budget.query_timeout, budget.cancel_timeout, - execute_query_with_max_rows_inner(&client, sql, max_rows), + execute_query_with_max_rows_inner(&client, sql, max_rows, prefer_text_protocol), ) .await } @@ -3164,7 +3236,7 @@ pub async fn execute_query_with_schema_and_max_rows( "[postgres][execute_with_schema:skip-search-path] total_ms={} reason=transaction-recovery", start.elapsed().as_millis() ); - return execute_query_with_max_rows_inner(&client, sql, max_rows).await; + return execute_query_with_max_rows_inner(&client, sql, max_rows, false).await; } let set_schema_start = Instant::now(); @@ -3182,7 +3254,7 @@ pub async fn execute_query_with_schema_and_max_rows( ); let query_start = Instant::now(); - let result = execute_query_with_max_rows_inner(&client, sql, max_rows).await; + let result = execute_query_with_max_rows_inner(&client, sql, max_rows, false).await; if result.is_ok() { clear_postgres_caches_after_ddl(pool, Some(&client), sql); } @@ -3205,6 +3277,7 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( cancel_token: Option, budget: DbOperationBudget, cancel_context: Option, + prefer_text_protocol: bool, ) -> Result { let start = Instant::now(); let checkout_start = Instant::now(); @@ -3227,7 +3300,7 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( cancel_token, budget.query_timeout, budget.cancel_timeout, - execute_query_with_max_rows_inner(&client, sql, max_rows), + execute_query_with_max_rows_inner(&client, sql, max_rows, prefer_text_protocol), ) .await; } @@ -3254,7 +3327,7 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( cancel_token, budget.query_timeout, budget.cancel_timeout, - execute_query_with_max_rows_inner(&client, sql, max_rows), + execute_query_with_max_rows_inner(&client, sql, max_rows, prefer_text_protocol), ) .await; if result.is_ok() { @@ -3493,12 +3566,17 @@ async fn execute_query_with_max_rows_inner( client: &deadpool_postgres::Client, sql: &str, max_rows: Option, + prefer_text_protocol: bool, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { - execute_select_query(client, sql, start, row_limit).await + if prefer_text_protocol { + execute_select_text(client, sql, start, row_limit, None).await + } else { + execute_select_query(client, sql, start, row_limit).await + } } else { let affected = client.execute(sql, &[]).await.map_err(pg_error_to_string)?; @@ -5721,6 +5799,61 @@ mod tests { assert!(!sql.contains("pg_get_function_identity_arguments")); } + #[test] + fn redshift_columns_sql_uses_simple_information_schema_metadata() { + let sql = redshift_columns_sql("tenant's", "orders"); + assert!(sql.contains("FROM information_schema.columns c")); + assert!(sql.contains("c.table_schema = 'tenant''s'")); + assert!(sql.contains("c.table_name = 'orders'")); + assert!(!sql.contains("pg_attribute")); + assert!(!sql.contains("pg_index")); + assert!(!sql.contains('$')); + } + + #[test] + fn redshift_columns_from_text_result_preserves_basic_metadata() { + let result = QueryResult { + columns: vec![ + "column_name".to_string(), + "full_type".to_string(), + "is_nullable".to_string(), + "column_default".to_string(), + "numeric_precision".to_string(), + "numeric_scale".to_string(), + "character_maximum_length".to_string(), + ], + column_types: Vec::new(), + column_sortables: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows: vec![vec![ + serde_json::json!("amount"), + serde_json::json!("numeric"), + serde_json::json!("NO"), + serde_json::Value::Null, + serde_json::json!(18), + serde_json::json!(2), + serde_json::Value::Null, + ]], + affected_rows: 0, + execution_time_ms: 0, + truncated: false, + session_id: None, + has_more: false, + elasticsearch_raw_body: None, + }; + + let columns = redshift_columns_from_query_result(result); + assert_eq!(columns.len(), 1); + assert_eq!(columns[0].name, "amount"); + assert_eq!(columns[0].data_type, "numeric"); + assert!(!columns[0].is_nullable); + assert_eq!(columns[0].numeric_precision, Some(18)); + assert_eq!(columns[0].numeric_scale, Some(2)); + assert_eq!(columns[0].character_maximum_length, None); + assert!(!columns[0].is_primary_key); + } + #[test] fn function_identity_arguments_probe_uses_pg_proc() { let sql = postgres_has_function_identity_arguments_sql(); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 0e4248e54..3354ea576 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -989,6 +989,10 @@ fn query_pool_database<'a>(database: &'a str, catalog: Option<&str>) -> Option<& } } +fn postgres_prefers_text_protocol(db_type: Option) -> bool { + db_type == Some(DatabaseType::Redshift) +} + pub async fn operation_budget_for_pool_key( state: &AppState, pool_key: &str, @@ -1143,6 +1147,7 @@ pub async fn do_execute( let p = p.clone(); let schema = schema.map(|s| s.to_string()); let max_rows = options.max_rows; + let prefer_text_protocol = postgres_prefers_text_protocol(pool_db_type); let cancel_context = state.get_postgres_cancel_context(pool_key).await; drop(connections); if let Some(schema) = schema { @@ -1154,6 +1159,7 @@ pub async fn do_execute( cancel_token, operation_budget.clone(), cancel_context, + prefer_text_protocol, ) .await } else { @@ -1164,6 +1170,7 @@ pub async fn do_execute( cancel_token, operation_budget.clone(), cancel_context, + prefer_text_protocol, ) .await } @@ -3661,6 +3668,13 @@ pub async fn rollback_manual_transaction(state: &AppState, txn_session_id: &str) #[cfg(test)] mod tests { use super::*; + + #[test] + fn redshift_queries_prefer_text_protocol() { + assert!(postgres_prefers_text_protocol(Some(DatabaseType::Redshift))); + assert!(!postgres_prefers_text_protocol(Some(DatabaseType::Postgres))); + assert!(!postgres_prefers_text_protocol(None)); + } #[cfg(unix)] use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec}; use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 58f5c27c6..480355eab 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -5334,6 +5334,11 @@ async fn get_columns_core_for_session_inner( PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::get_columns(p, schema, table).await.map(deduplicate_column_infos) } + PoolKind::Postgres(p) + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Redshift) => + { + db::postgres::get_redshift_columns(p, schema, table).await.map(deduplicate_column_infos) + } PoolKind::Postgres(p) => db::postgres::get_columns(p, schema, table).await.map(deduplicate_column_infos), PoolKind::Sqlite(p) => db::sqlite::get_columns(p, schema, table).await.map(deduplicate_column_infos), PoolKind::Rqlite(client) => { @@ -5487,6 +5492,11 @@ async fn list_indexes_core_for_session( PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { db::questdb::list_indexes(p, schema, table).await } + PoolKind::Postgres(_) + if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Redshift) => + { + Ok(vec![]) + } PoolKind::Postgres(p) => db::postgres::list_indexes(p, schema, table).await, PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, schema, table).await, PoolKind::Rqlite(client) => db::rqlite_driver::list_indexes(client, schema, table).await,