use anyhow::anyhow; use chrono::Utc; use std::{collections::HashMap, str::FromStr, sync::Arc, vec}; use windmill_parser::Arg; use futures::{future::BoxFuture, FutureExt, StreamExt}; use itertools::Itertools; use oracle::sql_type::{InnerValue, OracleType, ToSql}; use serde::{Deserialize, Serialize}; use serde_json::{json, value::RawValue, Value}; use windmill_common::{ error::{to_anyhow, Error}, s3_helpers::convert_json_line_stream, worker::{to_raw_value, Connection}, }; use windmill_queue::MiniPulledJob; use windmill_parser_sql::{ parse_db_resource, parse_oracledb_sig, parse_s3_mode, parse_sql_blocks, parse_sql_statement_named_params, }; use windmill_queue::CanceledBy; use crate::{ common::{ build_args_values, check_executor_binary_exists, s3_mode_args_to_worker_data, OccupancyMetrics, S3ModeWorkerData, }, handle_child::run_future_with_polling_update_job_poller, sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args }; use windmill_common::client::AuthedClient; #[derive(Deserialize)] struct OracleDatabase { user: String, password: String, database: String, } lazy_static::lazy_static! { static ref ORACLE_LIB_DIR: String = std::env::var("ORACLE_LIB_DIR").unwrap_or_else(|_| "/opt/oracle/23/lib".to_string()); } pub fn do_oracledb_inner<'a>( query: &str, params: Vec<(String, Box)>, conn: Arc>, column_order: Option<&'a mut Option>>, skip_collect: bool, s3: Option, ) -> windmill_common::error::Result>>> { let qw = query.trim_end_matches(';').to_string(); let result_f = async move { let param_names = parse_sql_statement_named_params(&qw, ':') .into_iter() .map(|x| x.into_bytes()) .collect_vec(); if skip_collect { tokio::task::spawn_blocking(move || { let c = conn.lock()?; let params2: Vec<(&str, &dyn ToSql)> = params .iter() .filter(|(k, _)| param_names.contains(&k.clone().into_bytes())) .map(|(key, val)| (key.as_str(), &**val as &dyn ToSql)) .collect(); let mut stmt = c.statement(&qw).build()?; match stmt.statement_type() { oracle::StatementType::Select => { stmt.query_named(¶ms2)?; } _ => { stmt.execute_named(¶ms2)?; c.commit()?; } } oracle::Result::Ok(()) }) .await .map_err(to_anyhow)? .map_err(to_anyhow)?; Ok(to_raw_value(&Value::Array(vec![]))) } else { // We use an mpsc because we need an async stream for s3 mode. However since everything is sync // in rust-oracle, I assumed that calling ResultSet::next() is blocking when it has to refetch. let (tx, rx) = tokio::sync::mpsc::channel::>(1000); let (column_order_oneshot_tx, column_order_oneshot_rx) = tokio::sync::oneshot::channel::>>(); let mut column_order_oneshot_tx = Some(column_order_oneshot_tx); let rows_stream = tokio_stream::wrappers::ReceiverStream::new(rx); tokio::task::spawn_blocking(move || { let result = (|| { let tx = tx.clone(); let params2: Vec<(&str, &dyn ToSql)> = params .iter() .filter(|(k, _)| param_names.contains(&k.clone().into_bytes())) .map(|(key, val)| (key.as_str(), &**val as &dyn ToSql)) .collect(); let c = conn.lock()?; let mut stmt = c.statement(&qw).build()?; match stmt.statement_type() { oracle::StatementType::Select => { let mut result_rows = stmt.query_named(¶ms2)?.enumerate(); while let Some((i, row)) = result_rows.next() { match row { Ok(row) => { // If first row, infer column order and send it to the channel if i == 0 { let col_order: Vec = row .column_info() .iter() .map(|x| x.name().to_string()) .collect::>(); let _ = column_order_oneshot_tx .take() .unwrap() .send(Some(col_order)); } // called in a spawn_blocking synchronous context, unwrap won't panic tx.blocking_send(Ok(convert_row_to_value(row))).unwrap() } Err(e) => { tx.blocking_send(Err(e)).unwrap(); break; } } } } _ => { stmt.execute_named(¶ms2)?; c.commit()?; } }; drop(column_order_oneshot_tx); Ok::<_, oracle::Error>(()) })(); match result { Ok(_) => {} Err(e) => tx.blocking_send(Err(e)).unwrap(), } // all instances of tx should be dropped here }); if let Ok(Some(col_order)) = column_order_oneshot_rx.await { if let Some(column_order) = column_order { *column_order = Some(col_order); } } if let Some(s3) = s3 { let stream = convert_json_line_stream(rows_stream.boxed(), s3.format).await?; s3.upload(stream.boxed()).await?; return Ok(to_raw_value(&s3.to_return_s3_obj())); } else { let rows: Vec<_> = rows_stream.collect().await; Ok(to_raw_value( &rows .into_iter() .collect::, _>>() .map_err(to_anyhow)? .into_iter() .collect::>(), )) } } }; Ok(result_f.boxed()) } fn convert_row_to_value(row: oracle::Row) -> serde_json::Value { let mut map = serde_json::Map::new(); for (key, value) in row.column_info().iter().zip(row.sql_values()) { map.insert( key.name().to_string(), convert_oracledb_value_to_json(value, key.oracle_type()), ); } serde_json::Value::Object(map) } fn conversion_error(r: Result) -> serde_json::Value { match r { Ok(v) => json!(v), Err(e) => json!(format!("Error converting value: {:?}", e)), } } fn convert_oracledb_value_to_json(v: &oracle::SqlValue, c: &OracleType) -> serde_json::Value { match v.as_inner_value() { Err(_) => serde_json::Value::Null, Ok(iv) => match iv { InnerValue::Int64(n) => json!(n), InnerValue::UInt64(n) => json!(n), InnerValue::Float(n) => json!(n), InnerValue::Double(n) => json!(n), InnerValue::Char(n) => json!(String::from_utf8_lossy(n)), InnerValue::Number(n) => json!(n), InnerValue::Boolean(n) => json!(n), InnerValue::Timestamp(_) => conversion_error(v.get::()), InnerValue::IntervalDS(_) => conversion_error(v.get::()), InnerValue::IntervalYM(_) => conversion_error(v.get::()), InnerValue::Clob(_) => json!("Unsupported type Clob"), InnerValue::Blob(_) => json!("Unsupported type Blob"), InnerValue::Rowid(_) => json!("Unsupported type Rowid"), InnerValue::Object(_) => json!("Unsuppported type Object"), InnerValue::Stmt(_) => json!("Unsupported type Stmt"), InnerValue::Raw(b) => { match c { OracleType::Varchar2(_) | OracleType::NVarchar2(_) | OracleType::Char(_) | OracleType::NChar(_) | OracleType::Json | OracleType::Xml => { let s = String::from_utf8_lossy(b); json!(s) } // OracleType::Rowid => todo!(), // OracleType::Raw(_) => todo!(), // // OracleType::BinaryFloat => todo!(), // OracleType::BinaryDouble => todo!(), // OracleType::Number(_, _) => todo!(), // OracleType::Float(_) => todo!(), // OracleType::Int64 => todo!(), // OracleType::UInt64 => todo!(), // OracleType::Long => todo!(), // OracleType::LongRaw => todo!(), OracleType::Date | OracleType::Timestamp(_) | OracleType::TimestampTZ(_) | OracleType::TimestampLTZ(_) | OracleType::IntervalDS(_, _) | OracleType::IntervalYM(_) => conversion_error(v.get::()), // OracleType::CLOB => todo!(), // OracleType::NCLOB => todo!(), // OracleType::BLOB => todo!(), // OracleType::BFILE => todo!(), // OracleType::RefCursor => todo!(), // OracleType::Boolean => todo!(), // OracleType::Object(_) => todo!(), _ => json!(format!("Unsupported type: {c:?}")), } } _ => json!(format!("Unsupported type: {c:?}")), }, } } fn get_statement_values( sig: Vec, job_args: &HashMap, args_to_skip: &Vec, ) -> (Vec<(String, Box)>, Vec) { let mut statement_values = vec![]; let mut errors = vec![]; for arg in &sig { if args_to_skip.contains(&arg.name) { continue; } let arg_t = arg.otyp.clone().unwrap_or_else(|| "text".to_string()); let arg_n = arg.name.clone(); let oracle_v: Box = match job_args .get(arg.name.as_str()) .unwrap_or_else(|| &json!(null)) { // Value::Null => todo!(), Value::Bool(b) => Box::new(*b), Value::String(s) if arg_t == "timestamp" || arg_t == "datetime" || arg_t == "date" || arg_t == "time" => { if let Ok(d) = chrono::DateTime::::from_str(s.as_str()) { Box::new(d) } else { Box::new(s.clone()) } } Value::String(s) => Box::new(s.clone()), Value::Number(n) if n.is_i64() && (arg_t == "int" || arg_t == "integer" || arg_t == "smallint" || arg_t == "bigint") => { Box::new(n.as_i64().unwrap()) } Value::Number(n) if n.is_f64() && arg_t == "float" => { Box::new(n.as_f64().unwrap() as f32) } Value::Number(n) if n.is_i64() && arg_t == "float" => { Box::new(n.as_i64().unwrap() as f32) } Value::Number(n) if n.is_u64() && arg_t == "uint" => Box::new(n.as_u64().unwrap()), Value::Number(n) if n.is_f64() && (arg_t == "real" || arg_t == "dec" || arg_t == "fixed") => { Box::new(n.as_f64().unwrap()) } Value::Number(n) if n.is_i64() && (arg_t == "real" || arg_t == "dec" || arg_t == "fixed") => { Box::new(n.as_i64().unwrap() as f64) } value @ _ => { errors.push(format!( "Unsupported type in query: {value:?} and signature {arg_t:?} for {arg_n}" )); continue; } }; statement_values.push((arg_n, oracle_v)); } (statement_values, errors) } pub async fn do_oracledb( job: &MiniPulledJob, client: &AuthedClient, query: &str, conn: &Connection, mem_peak: &mut i32, canceled_by: &mut Option, worker_name: &str, column_order: &mut Option>, occupancy_metrics: &mut OccupancyMetrics, ) -> windmill_common::error::Result> { check_executor_binary_exists( "the Oracle client lib", ORACLE_LIB_DIR.as_str(), "Oracle Database", )?; let job_args = build_args_values(job, client, conn).await?; let inline_db_res_path = parse_db_resource(&query); let s3 = parse_s3_mode(&query)?.map(|s3| s3_mode_args_to_worker_data(s3, client.clone(), job)); let db_arg = if let Some(inline_db_res_path) = inline_db_res_path { Some( client .get_resource_value_interpolated::( &inline_db_res_path, Some(job.id.to_string()), ) .await?, ) } else { job_args.get("database").cloned() }; let database = if let Some(db) = db_arg { serde_json::from_value::(db) .map_err(|e| Error::ExecutionErr(e.to_string()))? } else { return Err(Error::BadRequest("Missing database argument".to_string())); }; let annotations = windmill_common::worker::SqlAnnotations::parse(query); let sig = parse_oracledb_sig(query) .map_err(|x| Error::ExecutionErr(x.to_string()))? .args; let (query, args_to_skip) = sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args)?; let (statement_values, errors) = get_statement_values(sig.clone(), &job_args, &args_to_skip); if !errors.is_empty() { return Err(Error::ExecutionErr(errors.join("\n"))); } if !oracle::InitParams::is_initialized() { let _ = oracle::InitParams::new() .oracle_client_lib_dir(ORACLE_LIB_DIR.as_str()) .map_err(|e| anyhow!("Failed to initialize oracle client: {e}"))? .init(); } let oracle_conn = tokio::task::spawn_blocking(|| { oracle::Connection::connect(database.user, database.password, database.database) .map_err(|e| Error::ExecutionErr(e.to_string())) }) .await .map_err(to_anyhow)??; let conn_a = Arc::new(std::sync::Mutex::new(oracle_conn)); let queries = parse_sql_blocks(&query); let result_f = if queries.len() > 1 { let f = async { let mut res: Vec> = vec![]; for (i, q) in queries.iter().enumerate() { let (vals, _) = get_statement_values(sig.clone(), &job_args, &args_to_skip); let r = do_oracledb_inner( q, vals, conn_a.clone(), None, annotations.return_last_result && i < queries.len() - 1, s3.clone(), )? .await?; res.push(r); } if annotations.return_last_result && res.len() > 0 { Ok(res.pop().unwrap()) } else { Ok(to_raw_value(&res)) } }; f.boxed() } else { do_oracledb_inner( &query, statement_values, conn_a, Some(column_order), false, s3, )? }; let result = run_future_with_polling_update_job_poller( job.id, job.timeout, conn, mem_peak, canceled_by, result_f, worker_name, &job.workspace_id, &mut Some(occupancy_metrics), Box::pin(futures::stream::once(async { 0 })), ) .await?; let raw_result = windmill_common::worker::to_raw_value(&json!(result)); *mem_peak = (raw_result.get().len() / 1000) as i32; return Ok(raw_result); }