Expose logs

Makefile

@@ -207,7 +207,7 @@ validate:
 	cargo run --bin test-cli -- validate-datamodel dev_datamodel.prisma
 
 qe:
-	cargo run --bin query-engine -- --enable-playground --enable-raw-queries --enable-metrics --enable-open-telemetry
+	cargo run --bin query-engine -- --enable-playground --enable-raw-queries --enable-metrics --enable-open-telemetry --enable-logs-in-response
 
 qe-dmmf:
 	cargo run --bin query-engine -- cli dmmf > dmmf.json

query-engine/connector-test-kit-rs/query-tests-setup/src/runner/binary.rs

@@ -2,7 +2,8 @@ use crate::{ConnectorTag, RunnerInterface, TestError, TestResult, TxResult};
 use hyper::{Body, Method, Request, Response};
 use query_core::{schema::QuerySchemaRef, TxId};
 use query_engine::opt::PrismaOpt;
-use query_engine::server::{routes, setup, State};
+use query_engine::server::routes;
+use query_engine::state::{init_state, State};
 use query_engine_metrics::MetricRegistry;
 use request_handlers::{GQLBatchResponse, GQLError, GQLResponse, GraphQlBody, MultiQuery, PrismaResponse};
 
@@ -16,7 +17,7 @@ pub struct BinaryRunner {
 impl RunnerInterface for BinaryRunner {
     async fn load(datamodel: String, connector_tag: ConnectorTag, metrics: MetricRegistry) -> TestResult<Self> {
         let opts = PrismaOpt::from_list(&["binary", "--enable-raw-queries", "--datamodel", &datamodel]);
-        let state = setup(&opts, metrics).await.unwrap();
+        let state = init_state(&opts, false, Some(metrics)).await.unwrap();
 
         Ok(BinaryRunner {
             state,

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

@@ -1,6 +1,7 @@
 use once_cell::sync::Lazy;
 use opentelemetry::sdk::export::trace::SpanData;
-use opentelemetry::trace::TraceContextExt;
+use opentelemetry::trace::{TraceContextExt, TraceId};
+use opentelemetry::Context;
 use serde_json::{json, Value};
 use std::borrow::Cow;
 
@@ -95,6 +96,11 @@ pub fn set_span_link_from_trace_id(span: &Span, trace_id: Option<String>) {
     }
 }
 
+pub fn get_trace_id_from_context(context: &Context) -> TraceId {
+    let context_span = context.span();
+    context_span.span_context().trace_id()
+}
+
 pub fn is_user_facing_trace_filter(meta: &Metadata) -> bool {
     if !meta.is_span() {
         return false;

query-engine/query-engine/src/capture_tracer.rs

@@ -0,0 +1,108 @@
+use async_trait::async_trait;
+use opentelemetry::{
+    global,
+    sdk::{
+        self,
+        export::trace::{ExportResult, SpanData, SpanExporter},
+        propagation::TraceContextPropagator,
+    },
+    trace::{TraceId, TracerProvider},
+};
+use query_core::spans_to_json;
+use std::fmt::Debug;
+use std::{collections::HashMap, sync::Arc};
+use tokio::sync::Mutex;
+
+/// Pipeline builder
+#[derive(Debug)]
+pub struct PipelineBuilder {
+    trace_config: Option<sdk::trace::Config>,
+}
+
+/// Create a new stdout exporter pipeline builder.
+pub fn new_pipeline() -> PipelineBuilder {
+    PipelineBuilder::default()
+}
+
+impl Default for PipelineBuilder {
+    /// Return the default pipeline builder.
+    fn default() -> Self {
+        Self { trace_config: None }
+    }
+}
+
+impl PipelineBuilder {
+    /// Assign the SDK trace configuration.
+    pub fn with_trace_config(mut self, config: sdk::trace::Config) -> Self {
+        self.trace_config = Some(config);
+        self
+    }
+}
+
+impl PipelineBuilder {
+    pub fn install_simple(mut self, exporter: CaptureExporter) -> sdk::trace::Tracer {
+        global::set_text_map_propagator(TraceContextPropagator::new());
+
+        let mut provider_builder = sdk::trace::TracerProvider::builder().with_simple_exporter(exporter);
+        if let Some(config) = self.trace_config.take() {
+            provider_builder = provider_builder.with_config(config);
+        }
+        let provider = provider_builder.build();
+        let tracer = provider.tracer("opentelemetry");
+        global::set_tracer_provider(provider);
+
+        tracer
+    }
+}
+
+/// A [`CaptureExporter`] that sends spans to stdout.
+#[derive(Debug, Clone)]
+pub struct CaptureExporter {
+    logs: Arc<Mutex<HashMap<TraceId, Vec<SpanData>>>>,
+}
+
+impl CaptureExporter {
+    pub fn new() -> Self {
+        Self {
+            logs: Default::default(),
+        }
+    }
+
+    pub async fn capture(&self, trace_id: TraceId) {
+        let mut logs = self.logs.lock().await;
+        logs.insert(trace_id, Vec::new());
+    }
+
+    pub async fn get(&self, trace_id: TraceId) -> String {
+        let mut logs = self.logs.lock().await;
+        if let Some(spans) = logs.remove(&trace_id) {
+            spans_to_json(&spans)
+        } else {
+            String::new()
+        }
+    }
+}
+
+impl Default for CaptureExporter {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[async_trait]
+impl SpanExporter for CaptureExporter {
+    async fn export(&mut self, batch: Vec<SpanData>) -> ExportResult {
+        let batch = batch.into_iter().filter(|span| span.name == "quaint:query");
+
+        let mut logs = self.logs.lock().await;
+        for span in batch {
+            let trace_id = span.span_context.trace_id();
+
+            if let Some(spans) = logs.get_mut(&trace_id) {
+                spans.push(span)
+            }
+        }
+
+        Ok(())
+    }
+}

query-engine/query-engine/src/context.rs

@@ -1,7 +1,10 @@
-use crate::{PrismaError, PrismaResult};
+use crate::{capture_tracer::CaptureExporter, PrismaError, PrismaResult};
 use query_core::{executor, schema::QuerySchemaRef, schema_builder, QueryExecutor};
 use query_engine_metrics::MetricRegistry;
-use std::{env, fmt, sync::Arc};
+use std::{
+    env, fmt,
+    sync::{atomic::AtomicUsize, Arc},
+};
 
 /// Prisma request context containing all immutable state of the process.
 /// There is usually only one context initialized per process.
@@ -11,6 +14,8 @@ pub struct PrismaContext {
     pub metrics: MetricRegistry,
     /// Central query executor.
     pub executor: Box<dyn QueryExecutor + Send + Sync + 'static>,
+    pub inflight_tracer: Option<CaptureExporter>,
+    pub counter: AtomicUsize,
 }
 
 impl fmt::Debug for PrismaContext {
@@ -23,6 +28,7 @@ pub struct ContextBuilder {
     enable_raw_queries: bool,
     schema: psl::ValidatedSchema,
     metrics: Option<MetricRegistry>,
+    logs_capture_tracer: Option<CaptureExporter>,
 }
 
 impl ContextBuilder {
@@ -36,17 +42,29 @@ impl ContextBuilder {
         self
     }
 
+    pub fn set_logs_capture_tracer(mut self, tracer: Option<CaptureExporter>) -> Self {
+        self.logs_capture_tracer = tracer;
+        self
+    }
+
     pub async fn build(self) -> PrismaResult<PrismaContext> {
-        PrismaContext::new(self.schema, self.enable_raw_queries, self.metrics.unwrap_or_default()).await
+        PrismaContext::new(
+            self.schema,
+            self.enable_raw_queries,
+            self.metrics.unwrap_or_default(),
+            self.logs_capture_tracer,
+        )
+        .await
     }
 }
 
 impl PrismaContext {
     /// Initializes a new Prisma context.
-    async fn new(
+    pub async fn new(
         schema: psl::ValidatedSchema,
         enable_raw_queries: bool,
         metrics: MetricRegistry,
+        inflight_tracer: Option<CaptureExporter>,
     ) -> PrismaResult<Self> {
         let config = &schema.configuration;
         // We only support one data source at the moment, so take the first one (default not exposed yet).
@@ -70,6 +88,8 @@ impl PrismaContext {
             query_schema,
             executor,
             metrics,
+            inflight_tracer,
+            counter: AtomicUsize::new(0),
         };
 
         context.verify_connection().await?;
@@ -87,6 +107,7 @@ impl PrismaContext {
             enable_raw_queries: false,
             schema,
             metrics: None,
+            logs_capture_tracer: None,
         }
     }
 

query-engine/query-engine/src/lib.rs

@@ -1,11 +1,13 @@
 #![allow(clippy::derive_partial_eq_without_eq)]
 
+pub mod capture_tracer;
 pub mod cli;
 pub mod context;
 pub mod error;
 pub mod logger;
 pub mod opt;
 pub mod server;
+pub mod state;
 pub mod tracer;
 
 use error::PrismaError;

query-engine/query-engine/src/logger.rs

@@ -1,7 +1,5 @@
 use opentelemetry::{
-    global,
     sdk::{
-        propagation::TraceContextPropagator,
         trace::{Config, Tracer},
         Resource,
     },
@@ -13,7 +11,7 @@ use query_engine_metrics::MetricRegistry;
 use tracing::{dispatcher::SetGlobalDefaultError, subscriber};
 use tracing_subscriber::{filter::filter_fn, layer::SubscriberExt, EnvFilter, Layer};
 
-use crate::LogFormat;
+use crate::{capture_tracer::CaptureExporter, LogFormat};
 
 type LoggerResult<T> = Result<T, SetGlobalDefaultError>;
 
@@ -26,6 +24,7 @@ pub struct Logger<'a> {
     log_queries: bool,
     telemetry_endpoint: Option<&'a str>,
     metrics: Option<MetricRegistry>,
+    log_capture_exporter: Option<CaptureExporter>,
 }
 
 impl<'a> Logger<'a> {
@@ -38,6 +37,7 @@ impl<'a> Logger<'a> {
             log_queries: false,
             telemetry_endpoint: None,
             metrics: None,
+            log_capture_exporter: None,
         }
     }
 
@@ -69,6 +69,12 @@ impl<'a> Logger<'a> {
         self.metrics = Some(metrics);
     }
 
+    pub fn enable_logs_capture(&mut self) -> CaptureExporter {
+        let capture = CaptureExporter::new();
+        self.log_capture_exporter = Some(capture.clone());
+        capture
+    }
+
     /// Install logger as a global. Can be called only once per application
     /// instance. The returned guard value needs to stay in scope for the whole
     /// lifetime of the service.
@@ -79,7 +85,13 @@ impl<'a> Logger<'a> {
 
         let telemetry = if self.enable_telemetry {
             let tracer = create_otel_tracer(self.service_name, self.telemetry_endpoint);
-            let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+            let mut telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+
+            if let Some(log_capture_exporter) = self.log_capture_exporter {
+                let tracer = crate::capture_tracer::new_pipeline().install_simple(log_capture_exporter);
+                telemetry = telemetry.with_tracer(tracer);
+            }
+
             let telemetry = telemetry.with_filter(is_user_trace);
             Some(telemetry)
         } else {
@@ -109,8 +121,6 @@ impl<'a> Logger<'a> {
 }
 
 fn create_otel_tracer(service_name: &'static str, collector_endpoint: Option<&str>) -> Tracer {
-    global::set_text_map_propagator(TraceContextPropagator::new());
-
     if let Some(endpoint) = collector_endpoint {
         // A special parameter for Jaeger to set the service name in spans.
         let resource = Resource::new(vec![KeyValue::new("service.name", service_name)]);

query-engine/query-engine/src/main.rs

@@ -5,13 +5,13 @@ extern crate tracing;
 
 use query_engine::cli::CliCommand;
 use query_engine::error::PrismaError;
-use query_engine::logger::Logger;
 use query_engine::opt::PrismaOpt;
 use query_engine::server;
+use query_engine::state::init_state;
 use query_engine::LogFormat;
-use query_engine_metrics::MetricRegistry;
 use std::{error::Error, process};
 use structopt::StructOpt;
+use tracing::Instrument;
 
 type AnyError = Box<dyn Error + Send + Sync + 'static>;
 
@@ -26,26 +26,13 @@ async fn main() -> Result<(), AnyError> {
     async fn main() -> Result<(), PrismaError> {
         let opts = PrismaOpt::from_args();
 
-        let metrics = MetricRegistry::new();
-
-        let mut logger = Logger::new("prisma-engine-http");
-        logger.log_format(opts.log_format());
-        logger.log_queries(opts.log_queries());
-        logger.enable_telemetry(opts.enable_open_telemetry);
-        logger.telemetry_endpoint(&opts.open_telemetry_endpoint);
-        logger.enable_metrics(metrics.clone());
-
-        logger.install().unwrap();
-
-        if opts.enable_metrics || opts.dataproxy_metric_override {
-            query_engine_metrics::setup();
-        }
-
         match CliCommand::from_opt(&opts)? {
             Some(cmd) => cmd.execute().await?,
             None => {
+                let span = tracing::info_span!("prisma:engine:connect");
+                let state = init_state(&opts, true, None).instrument(span).await?;
                 set_panic_hook(opts.log_format());
-                server::listen(opts, metrics).await?;
+                server::listen(&opts, state).await?;
             }
         }
 

query-engine/query-engine/src/opt.rs

@@ -103,6 +103,10 @@ pub struct PrismaOpt {
     #[structopt(long)]
     pub enable_open_telemetry: bool,
 
+    #[structopt(long)]
+    /// Enable tracer to capture logs and return in the response
+    pub enable_logs_in_response: bool,
+
     /// The url to the OpenTelemetry collector.
     /// Enabling this will send the OpenTelemtry tracing to a collector
     /// and not via our custom stdout tracer