Repository: risingwavelabs/risingwave
File: src/utils/pgwire/src/pg_protocol.rs (Rust)

Code before:
    ///
    /// The span is only effective if there's a current session and the message is
    /// query-related. Otherwise, `Span::none()` is returned.
    fn root_span_for_msg(&self, msg: &FeMessage) -> tracing::Span {
        let Some(session_id) = self.session.as_ref().map(|s| s.id().0) else {
            return tracing::Span::none();
        };

        let mode = match msg {
            FeMessage::Query(_) => "simple query",
            FeMessage::Parse(_) => "extended query parse",
            FeMessage::Execute(_) => "extended query execute",
            _ => return tracing::Span::none(),
        };

        let mut span = tracing::info_span!(
            target: PGWIRE_ROOT_SPAN_TARGET,
            "handle_query",
            mode,
            session_id,
            sql = tracing::field::Empty,
        );
        if let Ok(sql) = msg.get_sql()
            && let Some(sql) = sql
        {
            record_sql_in_span(&sql, self.redact_sql_option_keywords.clone(), &mut span);
        }
        span
    }

    /// Return type `Option<()>` is essentially a bool, but allows `?` for early return.
    /// - `None` means to terminate the current connection
    /// - `Some(())` means to continue processing the next message
    async fn do_process(&mut self, msg: FeMessage) -> Option<()> {
        let span = self.root_span_for_msg(&msg);
        let weak_session = self
            .session
            .as_ref()
            .map(|s| Arc::downgrade(s) as Weak<dyn Any + Send + Sync>);

        // Processing the message itself.
        //
        // Note: pin the future to avoid stack overflow as we'll wrap it multiple times
        // in the following code.
        let fut = Box::pin(self.do_process_inner(msg));

        // Set the current session as the context when processing the message, if exists.
        let fut = async move {
            if let Some(session) = weak_session {
                CURRENT_SESSION.scope(session, fut).await
            } else {

Review comment (suggestion):
```suggestion
            record_sql_in_span(sql, self.redact_sql_option_keywords.clone(), &mut span);
```

Code after:
    ///
    /// The span is only effective if there's a current session and the message is
    /// query-related. Otherwise, `Span::none()` is returned.
    fn root_span_for_msg(&self, msg: &FeMessage) -> tracing::Span {
        let Some(session_id) = self.session.as_ref().map(|s| s.id().0) else {
            return tracing::Span::none();
        };

        let mode = match msg {
            FeMessage::Query(_) => "simple query",
            FeMessage::Parse(_) => "extended query parse",
            FeMessage::Execute(_) => "extended query execute",
            _ => return tracing::Span::none(),
        };

        let mut span = tracing::info_span!(
            target: PGWIRE_ROOT_SPAN_TARGET,
            "handle_query",
            mode,
            session_id,
            sql = tracing::field::Empty,
        );
        if let Ok(sql) = msg.get_sql()
            && let Some(sql) = sql
        {
            record_sql_in_span(sql, self.redact_sql_option_keywords.clone(), &mut span);
        }
        span
    }

    /// Return type `Option<()>` is essentially a bool, but allows `?` for early return.
    /// - `None` means to terminate the current connection
    /// - `Some(())` means to continue processing the next message
    async fn do_process(&mut self, msg: FeMessage) -> Option<()> {
        let span = self.root_span_for_msg(&msg);
        let weak_session = self
            .session
            .as_ref()
            .map(|s| Arc::downgrade(s) as Weak<dyn Any + Send + Sync>);

        // Processing the message itself.
        //
        // Note: pin the future to avoid stack overflow as we'll wrap it multiple times
        // in the following code.
        let fut = Box::pin(self.do_process_inner(msg));

        // Set the current session as the context when processing the message, if exists.
        let fut = async move {
            if let Some(session) = weak_session {
                CURRENT_SESSION.scope(session, fut).await
            } else {