refactor(backend-native): Extract with_session function

Author: mcheshkov

packages/cubejs-backend-native/src/cubesql_utils.rs

@@ -1,3 +1,4 @@
+use std::future::Future;
 use std::net::SocketAddr;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -49,3 +50,27 @@ pub async fn create_session(
 
     Ok(session)
 }
+
+pub async fn with_session<T, F, Fut>(
+    services: &NodeCubeServices,
+    native_auth_ctx: Arc<NativeAuthContext>,
+    f: F,
+) -> Result<T, CubeError>
+where
+    F: FnOnce(Arc<Session>) -> Fut,
+    Fut: Future<Output = Result<T, CubeError>>,
+{
+    let session_manager = services
+        .injector()
+        .get_service_typed::<SessionManager>()
+        .await;
+    let session = create_session(services, native_auth_ctx).await?;
+    let connection_id = session.state.connection_id;
+
+    // From now there's a session we should close before returning, as in `finally`
+    let result = { f(session).await };
+
+    session_manager.drop_session(connection_id).await;
+
+    result
+}

packages/cubejs-backend-native/src/node_export.rs

@@ -1,6 +1,5 @@
 use cubesql::compile::{convert_sql_to_cube_query, get_df_batches};
 use cubesql::config::processing_loop::ShutdownMode;
-use cubesql::sql::SessionManager;
 use cubesql::transport::TransportService;
 use futures::StreamExt;
 
@@ -11,7 +10,7 @@ use crate::auth::{NativeAuthContext, NodeBridgeAuthService};
 use crate::channel::call_js_fn;
 use crate::config::{NodeConfiguration, NodeConfigurationFactoryOptions, NodeCubeServices};
 use crate::cross::CLRepr;
-use crate::cubesql_utils::create_session;
+use crate::cubesql_utils::with_session;
 use crate::logger::NodeBridgeLogger;
 use crate::stream::OnDrainHandler;
 use crate::tokio_runtime_node;
@@ -185,77 +184,98 @@ async fn handle_sql_query(
 ) -> Result<(), CubeError> {
     let start_time = SystemTime::now();
 
-    let session = create_session(&services, native_auth_ctx.clone()).await?;
-
-    if let Some(auth_context) = session.state.auth_context() {
-        session
-            .session_manager
-            .server
-            .transport
-            .log_load_state(
-                None,
-                auth_context,
-                session.state.get_load_request_meta("sql"),
-                "Load Request".to_string(),
-                serde_json::json!({
-                    "query": {
-                        "sql": sql_query,
-                    }
-                }),
-            )
-            .await?;
-    }
-
     let transport_service = services
         .injector()
         .get_service_typed::<dyn TransportService>()
         .await;
-    let session_manager = services
-        .injector()
-        .get_service_typed::<SessionManager>()
-        .await;
 
-    let session_clone = Arc::clone(&session);
+    with_session(&services, native_auth_ctx.clone(), |session| async move {
+        if let Some(auth_context) = session.state.auth_context() {
+            session
+                .session_manager
+                .server
+                .transport
+                .log_load_state(
+                    None,
+                    auth_context,
+                    session.state.get_load_request_meta("sql"),
+                    "Load Request".to_string(),
+                    serde_json::json!({
+                        "query": {
+                            "sql": sql_query,
+                        }
+                    }),
+                )
+                .await?;
+        }
 
-    let execute = || async move {
-        // todo: can we use compiler_cache?
-        let meta_context = transport_service
-            .meta(native_auth_ctx)
-            .await
-            .map_err(|err| CubeError::internal(format!("Failed to get meta context: {}", err)))?;
-        let query_plan = convert_sql_to_cube_query(sql_query, meta_context, session).await?;
+        let session_clone = Arc::clone(&session);
 
-        let mut stream = get_df_batches(&query_plan).await?;
+        let execute = || async move {
+            // todo: can we use compiler_cache?
+            let meta_context = transport_service
+                .meta(native_auth_ctx)
+                .await
+                .map_err(|err| {
+                    CubeError::internal(format!("Failed to get meta context: {}", err))
+                })?;
+            let query_plan = convert_sql_to_cube_query(sql_query, meta_context, session).await?;
 
-        let semaphore = Arc::new(Semaphore::new(0));
+            let mut stream = get_df_batches(&query_plan).await?;
 
-        let drain_handler = OnDrainHandler::new(
-            channel.clone(),
-            stream_methods.stream.clone(),
-            semaphore.clone(),
-        );
+            let semaphore = Arc::new(Semaphore::new(0));
+
+            let drain_handler = OnDrainHandler::new(
+                channel.clone(),
+                stream_methods.stream.clone(),
+                semaphore.clone(),
+            );
+
+            drain_handler.handle(stream_methods.on.clone()).await?;
 
-        drain_handler.handle(stream_methods.on.clone()).await?;
+            let mut is_first_batch = true;
+            while let Some(batch) = stream.next().await {
+                let (columns, data) = batch_to_rows(batch?)?;
 
-        let mut is_first_batch = true;
-        while let Some(batch) = stream.next().await {
-            let (columns, data) = batch_to_rows(batch?)?;
+                if is_first_batch {
+                    let mut schema = Map::new();
+                    schema.insert("schema".into(), columns);
+                    let columns = format!(
+                        "{}{}",
+                        serde_json::to_string(&serde_json::Value::Object(schema))?,
+                        CHUNK_DELIM
+                    );
+                    is_first_batch = false;
+
+                    call_js_fn(
+                        channel.clone(),
+                        stream_methods.write.clone(),
+                        Box::new(|cx| {
+                            let arg = cx.string(columns).upcast::<JsValue>();
+
+                            Ok(vec![arg.upcast::<JsValue>()])
+                        }),
+                        Box::new(|cx, v| match v.downcast_or_throw::<JsBoolean, _>(cx) {
+                            Ok(v) => Ok(v.value(cx)),
+                            Err(_) => Err(CubeError::internal(
+                                "Failed to downcast write response".to_string(),
+                            )),
+                        }),
+                        stream_methods.stream.clone(),
+                    )
+                    .await?;
+                }
 
-            if is_first_batch {
-                let mut schema = Map::new();
-                schema.insert("schema".into(), columns);
-                let columns = format!(
-                    "{}{}",
-                    serde_json::to_string(&serde_json::Value::Object(schema))?,
-                    CHUNK_DELIM
-                );
-                is_first_batch = false;
+                let mut rows = Map::new();
+                rows.insert("data".into(), serde_json::Value::Array(data));
+                let data = format!("{}{}", serde_json::to_string(&rows)?, CHUNK_DELIM);
+                let js_stream_write_fn = stream_methods.write.clone();
 
-                call_js_fn(
+                let should_pause = !call_js_fn(
                     channel.clone(),
-                    stream_methods.write.clone(),
+                    js_stream_write_fn,
                     Box::new(|cx| {
-                        let arg = cx.string(columns).upcast::<JsValue>();
+                        let arg = cx.string(data).upcast::<JsValue>();
 
                         Ok(vec![arg.upcast::<JsValue>()])
                     }),
@@ -268,93 +288,67 @@ async fn handle_sql_query(
                     stream_methods.stream.clone(),
                 )
                 .await?;
-            }
-
-            let mut rows = Map::new();
-            rows.insert("data".into(), serde_json::Value::Array(data));
-            let data = format!("{}{}", serde_json::to_string(&rows)?, CHUNK_DELIM);
-            let js_stream_write_fn = stream_methods.write.clone();
 
-            let should_pause = !call_js_fn(
-                channel.clone(),
-                js_stream_write_fn,
-                Box::new(|cx| {
-                    let arg = cx.string(data).upcast::<JsValue>();
-
-                    Ok(vec![arg.upcast::<JsValue>()])
-                }),
-                Box::new(|cx, v| match v.downcast_or_throw::<JsBoolean, _>(cx) {
-                    Ok(v) => Ok(v.value(cx)),
-                    Err(_) => Err(CubeError::internal(
-                        "Failed to downcast write response".to_string(),
-                    )),
-                }),
-                stream_methods.stream.clone(),
-            )
-            .await?;
-
-            if should_pause {
-                let permit = semaphore.acquire().await?;
-                permit.forget();
+                if should_pause {
+                    let permit = semaphore.acquire().await?;
+                    permit.forget();
+                }
             }
-        }
 
-        Ok::<(), CubeError>(())
-    };
+            Ok::<(), CubeError>(())
+        };
 
-    let result = execute().await;
-    let duration = start_time.elapsed().unwrap().as_millis() as u64;
+        let result = execute().await;
+        let duration = start_time.elapsed().unwrap().as_millis() as u64;
 
-    match &result {
-        Ok(_) => {
-            session_clone
-                .session_manager
-                .server
-                .transport
-                .log_load_state(
-                    None,
-                    session_clone.state.auth_context().unwrap(),
-                    session_clone.state.get_load_request_meta("sql"),
-                    "Load Request Success".to_string(),
-                    serde_json::json!({
-                        "query": {
-                            "sql": sql_query,
-                        },
-                        "apiType": "sql",
-                        "duration": duration,
-                        "isDataQuery": true
-                    }),
-                )
-                .await?;
-        }
-        Err(err) => {
-            session_clone
-                .session_manager
-                .server
-                .transport
-                .log_load_state(
-                    None,
-                    session_clone.state.auth_context().unwrap(),
-                    session_clone.state.get_load_request_meta("sql"),
-                    "Cube SQL Error".to_string(),
-                    serde_json::json!({
-                        "query": {
-                            "sql": sql_query
-                        },
-                        "apiType": "sql",
-                        "duration": duration,
-                        "error": err.message,
-                    }),
-                )
-                .await?;
+        match &result {
+            Ok(_) => {
+                session_clone
+                    .session_manager
+                    .server
+                    .transport
+                    .log_load_state(
+                        None,
+                        session_clone.state.auth_context().unwrap(),
+                        session_clone.state.get_load_request_meta("sql"),
+                        "Load Request Success".to_string(),
+                        serde_json::json!({
+                            "query": {
+                                "sql": sql_query,
+                            },
+                            "apiType": "sql",
+                            "duration": duration,
+                            "isDataQuery": true
+                        }),
+                    )
+                    .await?;
+            }
+            Err(err) => {
+                session_clone
+                    .session_manager
+                    .server
+                    .transport
+                    .log_load_state(
+                        None,
+                        session_clone.state.auth_context().unwrap(),
+                        session_clone.state.get_load_request_meta("sql"),
+                        "Cube SQL Error".to_string(),
+                        serde_json::json!({
+                            "query": {
+                                "sql": sql_query
+                            },
+                            "apiType": "sql",
+                            "duration": duration,
+                            "error": err.message,
+                        }),
+                    )
+                    .await?;
+            }
         }
-    }
-
-    session_manager
-        .drop_session(session_clone.state.connection_id)
-        .await;
 
-    result
+        result
+    })
+    .await
 }
 
 struct WritableStreamMethods {