fix(qe) Add in dispatcher for self_contained queries (prisma#2826)

query-engine/core/src/executor/execute_operation.rs

@@ -2,6 +2,7 @@ use super::pipeline::QueryPipeline;
 use crate::{IrSerializer, Operation, QueryGraph, QueryGraphBuilder, QueryInterpreter, QuerySchemaRef, ResponseData};
 use connector::{Connection, ConnectionLike, Connector};
 use futures::future;
+use tracing_futures::WithSubscriber;
 
 pub async fn execute_single_operation(
     query_schema: QuerySchemaRef,
@@ -63,18 +64,16 @@ pub async fn execute_many_self_contained<C: Connector + Send + Sync>(
 ) -> crate::Result<Vec<crate::Result<ResponseData>>> {
     let mut futures = Vec::with_capacity(operations.len());
 
+    let dispatcher = crate::get_current_dispatcher();
     for op in operations {
         match QueryGraphBuilder::new(query_schema.clone()).build(op.clone()) {
             Ok((graph, serializer)) => {
                 let conn = connector.get_connection().await?;
 
-                futures.push(tokio::spawn(execute_self_contained(
-                    conn,
-                    graph,
-                    serializer,
-                    force_transactions,
-                    trace_id.clone(),
-                )));
+                futures.push(tokio::spawn(
+                    execute_self_contained(conn, graph, serializer, force_transactions, trace_id.clone())
+                        .with_subscriber(dispatcher.clone()),
+                ));
             }
 
             // This looks unnecessary, but is the simplest way to preserve ordering of results for the batch.

query-engine/core/src/executor/mod.rs

@@ -17,6 +17,8 @@ use crate::{query_document::Operation, response_ir::ResponseData, schema::QueryS
 use async_trait::async_trait;
 use connector::Connector;
 
+use tracing::Dispatch;
+
 #[async_trait]
 pub trait QueryExecutor: TransactionManager {
     /// Executes a single operation and returns its result.
@@ -68,3 +70,20 @@ pub trait TransactionManager {
     /// Rolls back a transaction.
     async fn rollback_tx(&self, tx_id: TxId) -> crate::Result<()>;
 }
+
+// With the node-api when a future is spawned in a new thread `tokio:spawn` it will not
+// use the current dispatcher and its logs will not be captured anymore. We can use this
+// method to get the current dispatcher and combine it with `with_subscriber`
+// let dispatcher = get_current_dispatcher();
+// tokio::spawn(async {
+//      my_async_ops.await
+// }.with_subscriber(dispatcher));
+//
+//
+// Finally, this can be replaced with with_current_collector
+// https://github.com/tokio-rs/tracing/blob/master/tracing-futures/src/lib.rs#L234
+// once this is in a release
+
+pub fn get_current_dispatcher() -> Dispatch {
+    tracing::dispatcher::get_default(|current| current.clone())
+}

query-engine/core/src/interactive_transactions/actors.rs

@@ -10,7 +10,6 @@ use tokio::{
     },
     time::{self, Duration},
 };
-use tracing::Dispatch;
 use tracing_futures::WithSubscriber;
 
 use super::{CachedTx, TransactionError, TxOpRequest, TxOpRequestMsg, TxOpResponse};
@@ -254,19 +253,6 @@ impl ITXClient {
     }
 }
 
-// Warning!!! This is not pretty. When the actors are spawned with the node-api
-// They lose the current subscriber because it is not a global subscriber
-// and are no longer able to send logs. This fix gets the current dispatcher
-// and clones it out so that we can use `WithSubscriber` with it.
-fn get_current_dispatcher() -> Dispatch {
-    let mut dispatcher = tracing::Dispatch::default();
-    tracing::dispatcher::get_default(|current| {
-        dispatcher = current.clone();
-    });
-
-    dispatcher
-}
-
 pub fn spawn_itx_actor(
     query_schema: QuerySchemaRef,
     tx_id: TxId,
@@ -284,7 +270,7 @@ pub fn spawn_itx_actor(
     };
 
     let mut server = ITXServer::new(tx_id, CachedTx::Open(value), timeout, rx_from_client, query_schema);
-    let dispatcher = get_current_dispatcher();
+    let dispatcher = crate::get_current_dispatcher();
 
     tokio::task::spawn(
         async move {