|
| 1 | +use std::sync::Arc; |
| 2 | + |
| 3 | +use neon::prelude::*; |
| 4 | + |
| 5 | +use cubesql::compile::convert_sql_to_cube_query; |
| 6 | +use cubesql::compile::datafusion::logical_plan::LogicalPlan; |
| 7 | +use cubesql::compile::engine::df::scan::CubeScanNode; |
| 8 | +use cubesql::compile::engine::df::wrapper::{CubeScanWrappedSqlNode, CubeScanWrapperNode}; |
| 9 | +use cubesql::sql::Session; |
| 10 | +use cubesql::transport::MetaContext; |
| 11 | +use cubesql::CubeError; |
| 12 | + |
| 13 | +use crate::auth::NativeAuthContext; |
| 14 | +use crate::config::NodeCubeServices; |
| 15 | +use crate::cubesql_utils::with_session; |
| 16 | +use crate::tokio_runtime_node; |
| 17 | + |
| 18 | +enum Sql4SqlQueryType { |
| 19 | + Regular, |
| 20 | + PostProcessing, |
| 21 | + Pushdown, |
| 22 | +} |
| 23 | + |
| 24 | +impl Sql4SqlQueryType { |
| 25 | + pub fn to_js<'ctx>(&self, cx: &mut impl Context<'ctx>) -> JsResult<'ctx, JsString> { |
| 26 | + let self_str = match self { |
| 27 | + Self::Regular => "regular", |
| 28 | + Self::PostProcessing => "post_processing", |
| 29 | + Self::Pushdown => "pushdown", |
| 30 | + }; |
| 31 | + |
| 32 | + Ok(cx.string(self_str)) |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +enum Sql4SqlResponseResult { |
| 37 | + Ok { |
| 38 | + sql: String, |
| 39 | + values: Vec<Option<String>>, |
| 40 | + }, |
| 41 | + Error { |
| 42 | + error: String, |
| 43 | + }, |
| 44 | +} |
| 45 | + |
| 46 | +struct Sql4SqlResponse { |
| 47 | + result: Sql4SqlResponseResult, |
| 48 | + query_type: Sql4SqlQueryType, |
| 49 | +} |
| 50 | + |
| 51 | +impl Sql4SqlResponse { |
| 52 | + pub fn to_js<'ctx>(&self, cx: &mut impl Context<'ctx>) -> JsResult<'ctx, JsObject> { |
| 53 | + let obj = cx.empty_object(); |
| 54 | + |
| 55 | + match &self.result { |
| 56 | + Sql4SqlResponseResult::Ok { sql, values } => { |
| 57 | + let status = cx.string("ok"); |
| 58 | + obj.set(cx, "status", status)?; |
| 59 | + |
| 60 | + let sql_tuple = cx.empty_array(); |
| 61 | + let sql = cx.string(sql); |
| 62 | + sql_tuple.set(cx, 0, sql)?; |
| 63 | + let js_values = cx.empty_array(); |
| 64 | + for (i, v) in values.iter().enumerate() { |
| 65 | + use std::convert::TryFrom; |
| 66 | + let i = u32::try_from(i).unwrap(); |
| 67 | + let v: Handle<JsValue> = v |
| 68 | + .as_ref() |
| 69 | + .map(|v| cx.string(v).upcast()) |
| 70 | + .unwrap_or_else(|| cx.null().upcast()); |
| 71 | + js_values.set(cx, i, v)?; |
| 72 | + } |
| 73 | + sql_tuple.set(cx, 1, js_values)?; |
| 74 | + obj.set(cx, "sql", sql_tuple)?; |
| 75 | + } |
| 76 | + Sql4SqlResponseResult::Error { error } => { |
| 77 | + let status = cx.string("error"); |
| 78 | + obj.set(cx, "status", status)?; |
| 79 | + |
| 80 | + let error = cx.string(error); |
| 81 | + obj.set(cx, "error", error)?; |
| 82 | + } |
| 83 | + } |
| 84 | + |
| 85 | + let query_type = self.query_type.to_js(cx)?; |
| 86 | + obj.set(cx, "query_type", query_type)?; |
| 87 | + |
| 88 | + Ok(obj) |
| 89 | + } |
| 90 | +} |
| 91 | + |
| 92 | +async fn get_sql( |
| 93 | + session: &Session, |
| 94 | + meta_context: Arc<MetaContext>, |
| 95 | + plan: Arc<LogicalPlan>, |
| 96 | +) -> Result<Sql4SqlResponse, CubeError> { |
| 97 | + let auth_context = session |
| 98 | + .state |
| 99 | + .auth_context() |
| 100 | + .ok_or_else(|| CubeError::internal("Unexpected missing auth context".to_string()))?; |
| 101 | + |
| 102 | + match plan.as_ref() { |
| 103 | + LogicalPlan::Extension(extension) => { |
| 104 | + let cube_scan_wrapped_sql = extension |
| 105 | + .node |
| 106 | + .as_any() |
| 107 | + .downcast_ref::<CubeScanWrappedSqlNode>(); |
| 108 | + |
| 109 | + if let Some(cube_scan_wrapped_sql) = cube_scan_wrapped_sql { |
| 110 | + return Ok(Sql4SqlResponse { |
| 111 | + result: Sql4SqlResponseResult::Ok { |
| 112 | + sql: cube_scan_wrapped_sql.wrapped_sql.sql.clone(), |
| 113 | + values: cube_scan_wrapped_sql.wrapped_sql.values.clone(), |
| 114 | + }, |
| 115 | + query_type: Sql4SqlQueryType::Pushdown, |
| 116 | + }); |
| 117 | + } |
| 118 | + |
| 119 | + if extension.node.as_any().is::<CubeScanNode>() { |
| 120 | + let cube_scan_wrapper = CubeScanWrapperNode::new( |
| 121 | + plan, |
| 122 | + meta_context, |
| 123 | + auth_context, |
| 124 | + None, |
| 125 | + session.server.config_obj.clone(), |
| 126 | + ); |
| 127 | + let wrapped_sql = cube_scan_wrapper |
| 128 | + .generate_sql( |
| 129 | + session.server.transport.clone(), |
| 130 | + Arc::new(session.state.get_load_request_meta("sql")), |
| 131 | + ) |
| 132 | + .await?; |
| 133 | + |
| 134 | + return Ok(Sql4SqlResponse { |
| 135 | + result: Sql4SqlResponseResult::Ok { |
| 136 | + sql: wrapped_sql.wrapped_sql.sql.clone(), |
| 137 | + values: wrapped_sql.wrapped_sql.values.clone(), |
| 138 | + }, |
| 139 | + query_type: Sql4SqlQueryType::Regular, |
| 140 | + }); |
| 141 | + } |
| 142 | + |
| 143 | + Err(CubeError::internal( |
| 144 | + "Unexpected extension in logical plan root".to_string(), |
| 145 | + )) |
| 146 | + } |
| 147 | + _ => Ok(Sql4SqlResponse { |
| 148 | + result: Sql4SqlResponseResult::Error { |
| 149 | + error: "Provided query can not be executed without post-processing.".to_string(), |
| 150 | + }, |
| 151 | + query_type: Sql4SqlQueryType::PostProcessing, |
| 152 | + }), |
| 153 | + } |
| 154 | +} |
| 155 | + |
| 156 | +async fn handle_sql4sql_query( |
| 157 | + services: Arc<NodeCubeServices>, |
| 158 | + native_auth_ctx: Arc<NativeAuthContext>, |
| 159 | + sql_query: &str, |
| 160 | +) -> Result<Sql4SqlResponse, CubeError> { |
| 161 | + with_session(&services, native_auth_ctx.clone(), |session| async move { |
| 162 | + let transport = session.server.transport.clone(); |
| 163 | + // todo: can we use compiler_cache? |
| 164 | + let meta_context = transport |
| 165 | + .meta(native_auth_ctx) |
| 166 | + .await |
| 167 | + .map_err(|err| CubeError::internal(format!("Failed to get meta context: {err}")))?; |
| 168 | + let query_plan = |
| 169 | + convert_sql_to_cube_query(sql_query, meta_context.clone(), session.clone()).await?; |
| 170 | + let logical_plan = query_plan.try_as_logical_plan()?; |
| 171 | + get_sql(&session, meta_context, Arc::new(logical_plan.clone())).await |
| 172 | + }) |
| 173 | + .await |
| 174 | +} |
| 175 | + |
| 176 | +pub fn sql4sql(mut cx: FunctionContext) -> JsResult<JsValue> { |
| 177 | + let interface = cx.argument::<JsBox<crate::node_export::SQLInterface>>(0)?; |
| 178 | + let sql_query = cx.argument::<JsString>(1)?.value(&mut cx); |
| 179 | + |
| 180 | + let security_context: Option<serde_json::Value> = match cx.argument::<JsValue>(2) { |
| 181 | + Ok(string) => match string.downcast::<JsString, _>(&mut cx) { |
| 182 | + Ok(v) => v.value(&mut cx).parse::<serde_json::Value>().ok(), |
| 183 | + Err(_) => None, |
| 184 | + }, |
| 185 | + Err(_) => None, |
| 186 | + }; |
| 187 | + |
| 188 | + let services = interface.services.clone(); |
| 189 | + let runtime = tokio_runtime_node(&mut cx)?; |
| 190 | + |
| 191 | + let channel = cx.channel(); |
| 192 | + |
| 193 | + let native_auth_ctx = Arc::new(NativeAuthContext { |
| 194 | + user: Some(String::from("unknown")), |
| 195 | + superuser: false, |
| 196 | + security_context, |
| 197 | + }); |
| 198 | + |
| 199 | + let (deferred, promise) = cx.promise(); |
| 200 | + |
| 201 | + // In case spawned task panics or gets aborted before settle call it will leave permanently pending Promise in JS land |
| 202 | + // We don't want to just waste whole thread (doesn't really matter main or worker or libuv thread pool) |
| 203 | + // just busy waiting that JoinHandle |
| 204 | + // TODO handle JoinError |
| 205 | + // keep JoinHandle alive in JS thread |
| 206 | + // check join handle from JS thread periodically, reject promise on JoinError |
| 207 | + // maybe register something like uv_check handle (libuv itself does not have ABI stability of N-API) |
| 208 | + // can do it relatively rare, and in a single loop for all JoinHandles |
| 209 | + // this is just a watchdog for a Very Bad case, so latency requirement can be quite relaxed |
| 210 | + runtime.spawn(async move { |
| 211 | + let result = handle_sql4sql_query(services, native_auth_ctx, &sql_query).await; |
| 212 | + |
| 213 | + if let Err(err) = deferred.try_settle_with(&channel, move |mut cx| { |
| 214 | + // `neon::result::ResultExt` is implemented only for Result<Handle, Handle>, even though Ok variant is not touched |
| 215 | + let response = result.or_else(|err| cx.throw_error(err.to_string()))?; |
| 216 | + let response = response.to_js(&mut cx)?; |
| 217 | + Ok(response) |
| 218 | + }) { |
| 219 | + // There is not much we can do at this point |
| 220 | + // TODO lift this error to task => JoinHandle => JS watchdog |
| 221 | + log::error!( |
| 222 | + "Unable to settle JS promise from tokio task, try_settle_with failed, err: {err}" |
| 223 | + ); |
| 224 | + } |
| 225 | + }); |
| 226 | + |
| 227 | + Ok(promise.upcast::<JsValue>()) |
| 228 | +} |
0 commit comments