omplement get_final_cubestore_result & get_final_cubestore_result_multi in native

Author: KSDaemon

packages/cubejs-api-gateway/package.json

@@ -32,6 +32,7 @@
     "@ungap/structured-clone": "^0.3.4",
     "body-parser": "^1.19.0",
     "chrono-node": "^2.6.2",
+    "express": "^4.17.1",
     "express-graphql": "^0.12.0",
     "graphql": "^15.8.0",
     "graphql-scalars": "^1.10.0",
@@ -59,7 +60,6 @@
     "@types/ramda": "^0.27.32",
     "@types/supertest": "^2.0.10",
     "@types/uuid": "^8.3.1",
-    "express": "^4.17.1",
     "jest": "^27",
     "mysql": "^2.18.1",
     "should": "^13.2.3",

packages/cubejs-api-gateway/src/gateway.ts

@@ -10,7 +10,11 @@ import {
   getRealType,
   QueryAlias,
 } from '@cubejs-backend/shared';
-import { transformData as transformDataNative, TransformDataResponse } from '@cubejs-backend/native';
+import {
+  getFinalCubestoreResult, getFinalCubestoreResultMulti,
+  transformData as transformDataNative,
+  TransformDataResponse
+} from '@cubejs-backend/native';
 import type {
   Application as ExpressApplication,
   ErrorRequestHandler,
@@ -114,6 +118,16 @@ function systemAsyncHandler(handler: (req: Request & { context: ExtendedRequestC
   };
 }
 
+function cleanupResult(result) {
+  return {
+    ...result,
+    dataCb: undefined,
+    rawData: undefined,
+    transformDataParams: undefined,
+    isNative: undefined,
+  };
+}
+
 // Prepared CheckAuthFn, default or from config: always async, returns nothing
 type PreparedCheckAuthFn = (ctx: any, authorization?: string) => Promise<void>;
 
@@ -1623,6 +1637,7 @@ class ApiGateway {
     return {
       query: normalizedQuery,
       dataCb,
+      rawData: response.data,
       transformDataParams,
       lastRefreshTime: response.lastRefreshTime?.toISOString(),
       ...(
@@ -1805,27 +1820,56 @@ class ApiGateway {
       const allNative = results.every(r => r.isNative);
 
       if (props.queryType === 'multi') {
-        res({
-          queryType,
-          results: results.map(r => {
-            const data = r.dataCb();
-            return {
-              ...r,
-              data,
-              dataCb: undefined,
-              transformDataParams: undefined
-            };
-          }),
-          pivotQuery: getPivotQuery(queryType, normalizedQueries),
-          slowQuery
-        });
+        // If all query results are from Cubestore (are native)
+        // we prepare the final json result on native side
+        if (allNative) {
+          const [transformDataJson, rawDataRef, cleanResultList] = results.reduce<[string[], any[], string[]]>(
+            ([transformList, rawList, resultList], r) => {
+              transformList.push(JSON.stringify(r.transformDataParams));
+              rawList.push(r.rawData.getNativeRef());
+              resultList.push(cleanupResult(r));
+              return [transformList, rawList, resultList];
+            },
+            [[], [], []]
+          );
+
+          const responseDataObj = {
+            queryType,
+            results: cleanResultList,
+            slowQuery
+          };
+          const resultDataJson = JSON.stringify(responseDataObj);
+
+          res(getFinalCubestoreResultMulti(transformDataJson, rawDataRef, resultDataJson));
+        } else {
+          // if we have mixed query results (there are js and native)
+          // we prepare results separately: on js and native sides
+          // and serve final response from JS side
+          res({
+            queryType,
+            results: results.map(r => {
+              const data = r.dataCb();
+              return {
+                ...cleanupResult(r),
+                data,
+              };
+            }),
+            pivotQuery: getPivotQuery(queryType, normalizedQueries),
+            slowQuery
+          });
+        }
+      } else if (allNative) {
+        // We prepare the full final json result on native side
+        const r = results[0];
+        const transformDataJson = JSON.stringify(r.transformDataParams);
+        const rawDataRef = r.rawData.getNativeRef();
+        const resultDataJson = JSON.stringify(cleanupResult(r));
+        res(getFinalCubestoreResult(transformDataJson, rawDataRef, resultDataJson));
       } else {
         const data = results[0].dataCb();
         res({
-          ...results[0],
+          ...cleanupResult(results[0]),
           data,
-          dataCb: undefined,
-          transformDataParams: undefined
         });
       }
     } catch (e: any) {
@@ -2034,10 +2078,10 @@ class ApiGateway {
 
       if (message instanceof ArrayBuffer) {
         res.set('Content-Type', 'application/json');
-        res.send(message);
+        res.send(Buffer.from(message));
+      } else {
+        res.json(message);
       }
-
-      res.json(message);
     };
   }
 

packages/cubejs-backend-native/Cargo.lock

@@ -389,6 +389,18 @@ export const transformData = (transformDataJson: string, rows: any): TransformDa
   return native.transformQueryData(transformDataJson, rows);
 };
 
+export const getFinalCubestoreResult = (transformDataJson: string, rows: any, resultData: string): ArrayBuffer => {
+  const native = loadNative();
+
+  return native.getFinalCubestoreResult(transformDataJson, rows, resultData);
+};
+
+export const getFinalCubestoreResultMulti = (transformDataJson: string[], rows: any[], responseData: string): ArrayBuffer => {
+  const native = loadNative();
+
+  return native.getFinalCubestoreResultMulti(transformDataJson, rows, responseData);
+};
+
 export interface PyConfiguration {
   repositoryFactory?: (ctx: unknown) => Promise<unknown>,
   logger?: (msg: string, params: Record<string, any>) => void,

packages/cubejs-backend-native/js/index.ts

@@ -36,10 +36,12 @@ use cubeorchestrator::cubestore_message_parser::CubeStoreResult;
 
 use cubesql::{telemetry::ReportingLogger, CubeError};
 
+use cubeorchestrator::cubestore_result_transform::{
+    get_final_cubestore_result, get_final_cubestore_result_multi, transform_data,
+};
+use cubeorchestrator::types::{RequestResultData, RequestResultDataMulti, TransformDataRequest};
 use neon::prelude::*;
 use neon::types::buffer::TypedArray;
-use cubeorchestrator::cubestore_result_transform::transform_data;
-use cubeorchestrator::types::{TransformDataRequest};
 
 struct SQLInterface {
     services: Arc<NodeCubeServices>,
@@ -546,9 +548,7 @@ fn get_cubestore_result(mut cx: FunctionContext) -> JsResult<JsValue> {
 }
 
 fn transform_query_data(mut cx: FunctionContext) -> JsResult<JsObject> {
-    let json_str = cx
-        .argument::<JsString>(0)?
-        .value(&mut cx);
+    let json_str = cx.argument::<JsString>(0)?.value(&mut cx);
     let request_data = match serde_json::from_str::<TransformDataRequest>(&json_str) {
         Ok(data) => data,
         Err(err) => return cx.throw_error(err.to_string()),
@@ -571,7 +571,7 @@ fn transform_query_data(mut cx: FunctionContext) -> JsResult<JsObject> {
         res_type.clone(),
     ) {
         Ok(data) => data,
-        Err(err) => return cx.throw_error(err.to_string())
+        Err(err) => return cx.throw_error(err.to_string()),
     };
 
     let json_data = match serde_json::to_string(&transformed) {
@@ -587,6 +587,98 @@ fn transform_query_data(mut cx: FunctionContext) -> JsResult<JsObject> {
     Ok(js_result)
 }
 
+fn final_cubestore_result(mut cx: FunctionContext) -> JsResult<JsArrayBuffer> {
+    let transform_data_str = cx.argument::<JsString>(0)?.value(&mut cx);
+    let transform_request_data =
+        match serde_json::from_str::<TransformDataRequest>(&transform_data_str) {
+            Ok(data) => data,
+            Err(err) => return cx.throw_error(err.to_string()),
+        };
+    let cube_store_result = cx.argument::<JsBox<CubeStoreResult>>(1)?;
+    let result_data_str = cx.argument::<JsString>(2)?.value(&mut cx);
+    let mut result_data = match serde_json::from_str::<RequestResultData>(&result_data_str) {
+        Ok(data) => data,
+        Err(err) => return cx.throw_error(err.to_string()),
+    };
+
+    if let Err(err) = get_final_cubestore_result(
+        &transform_request_data,
+        &**cube_store_result,
+        &mut result_data,
+    ) {
+        return cx.throw_error(err.to_string());
+    }
+
+    let json_data = match serde_json::to_string(&result_data) {
+        Ok(data) => data,
+        Err(e) => return cx.throw_error(format!("Serialization error: {}", e)),
+    };
+    let json_bytes = json_data.as_bytes();
+
+    let mut js_buffer = cx.array_buffer(json_bytes.len())?;
+    {
+        let buffer = js_buffer.as_mut_slice(&mut cx);
+        buffer.copy_from_slice(json_bytes);
+    }
+
+    Ok(js_buffer)
+}
+
+fn final_cubestore_result_multi(mut cx: FunctionContext) -> JsResult<JsArrayBuffer> {
+    let transform_data_array = cx.argument::<JsArray>(0)?;
+    let transform_requests: Vec<TransformDataRequest> = transform_data_array
+        .to_vec(&mut cx)?
+        .into_iter()
+        .map(|js_value| {
+            let js_string = js_value
+                .downcast_or_throw::<JsString, _>(&mut cx)?
+                .value(&mut cx);
+
+            match serde_json::from_str::<TransformDataRequest>(&js_string) {
+                Ok(request) => Ok(request),
+                Err(err) => return cx.throw_error(err.to_string()),
+            }
+        })
+        .collect::<Result<_, _>>()?;
+
+    let cube_store_array = cx.argument::<JsArray>(1)?;
+    let cube_store_results_boxed: Vec<Handle<JsBox<CubeStoreResult>>> = cube_store_array
+        .to_vec(&mut cx)?
+        .into_iter()
+        .map(|js_value| js_value.downcast_or_throw::<JsBox<CubeStoreResult>, _>(&mut cx))
+        .collect::<Result<_, _>>()?;
+    let cube_store_results: Vec<&CubeStoreResult> = cube_store_results_boxed
+        .iter()
+        .map(|handle| &***handle)
+        .collect();
+
+    let result_data_str = cx.argument::<JsString>(2)?.value(&mut cx);
+    let mut result_data = match serde_json::from_str::<RequestResultDataMulti>(&result_data_str) {
+        Ok(data) => data,
+        Err(err) => return cx.throw_error(err.to_string()),
+    };
+
+    if let Err(err) =
+        get_final_cubestore_result_multi(&transform_requests, &cube_store_results, &mut result_data)
+    {
+        return cx.throw_error(err.to_string());
+    }
+
+    let json_data = match serde_json::to_string(&result_data) {
+        Ok(data) => data,
+        Err(e) => return cx.throw_error(format!("Serialization error: {}", e)),
+    };
+    let json_bytes = json_data.as_bytes();
+
+    let mut js_buffer = cx.array_buffer(json_bytes.len())?;
+    {
+        let buffer = js_buffer.as_mut_slice(&mut cx);
+        buffer.copy_from_slice(json_bytes);
+    }
+
+    Ok(js_buffer)
+}
+
 pub fn register_module_exports<C: NodeConfiguration + 'static>(
     mut cx: ModuleContext,
 ) -> NeonResult<()> {
@@ -607,6 +699,8 @@ pub fn register_module_exports<C: NodeConfiguration + 'static>(
     )?;
     cx.export_function("getCubestoreResult", get_cubestore_result)?;
     cx.export_function("transformQueryData", transform_query_data)?;
+    cx.export_function("getFinalCubestoreResult", final_cubestore_result)?;
+    cx.export_function("getFinalCubestoreResultMulti", final_cubestore_result_multi)?;
 
     crate::template::template_register_module(&mut cx)?;