From ec6dcbdd64be64d0747cf3a4c3a3195400ca7797 Mon Sep 17 00:00:00 2001
From: benibus <bpharks@gmx.com>
Date: Tue, 12 Mar 2024 13:00:04 -0400
Subject: [PATCH] Add stream interface

---
 cpp/src/arrow/telemetry/logging.cc        | 34 +++++++++++++++++------
 cpp/src/arrow/telemetry/logging.h         | 21 +++++++++++++-
 cpp/src/arrow/telemetry/telemetry_test.cc |  3 ++
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/cpp/src/arrow/telemetry/logging.cc b/cpp/src/arrow/telemetry/logging.cc
index 320ab09c106a7..f8b9b8a4ee3bc 100644
--- a/cpp/src/arrow/telemetry/logging.cc
+++ b/cpp/src/arrow/telemetry/logging.cc
@@ -243,9 +243,6 @@ class OtelLogger : public Logger {
       return;
     }
 
-    const auto timestamp =
-        otel::common::SystemTimestamp(std::chrono::system_clock::now());
-
     auto log = logger_->CreateLogRecord();
     if (log == nullptr) {
       return;
@@ -262,7 +259,7 @@ class OtelLogger : public Logger {
 
     // We set the remaining attributes AFTER the custom attributes in AttributeHolder
     // because, in the event of key collisions, these should take precedence.
-    log->SetTimestamp(timestamp);
+    log->SetTimestamp(otel::common::SystemTimestamp(desc.timestamp));
     log->SetSeverity(ToOtelSeverity(desc.severity));
 
     auto span_ctx = otel::trace::Tracer::GetCurrentSpan()->GetContext();
@@ -270,10 +267,7 @@ class OtelLogger : public Logger {
     log->SetTraceId(span_ctx.trace_id());
     log->SetTraceFlags(span_ctx.trace_flags());
 
-    if (desc.body) {
-      auto body = *desc.body;
-      log->SetBody(ToOtel(body));
-    }
+    log->SetBody(ToOtel(desc.body));
 
     if (const auto& event = desc.event_id; event.is_valid()) {
       log->SetEventId(event.id, ToOtel(event.name));
@@ -343,5 +337,29 @@ std::unique_ptr<Logger> GlobalLogger::logger_ = nullptr;
 
 std::unique_ptr<Logger> MakeNoopLogger() { return std::make_unique<NoopLogger>(); }
 
+class LogMessage::Impl {
+ public:
+  Impl(LogLevel severity, Logger* logger) : logger(logger), severity(severity) {}
+
+  ~Impl() {
+    if (logger) {
+      auto body = stream.str();
+      LogDescriptor desc;
+      desc.body = body;
+      desc.severity = severity;
+      logger->Log(desc);
+    }
+  }
+
+  Logger* logger;
+  LogLevel severity;
+  std::stringstream stream;
+};
+
+LogMessage::LogMessage(LogLevel severity, Logger* logger)
+    : impl_(std::make_shared<Impl>(severity, logger)) {}
+
+std::ostream& LogMessage::Stream() { return impl_->stream; }
+
 }  // namespace telemetry
 }  // namespace arrow
diff --git a/cpp/src/arrow/telemetry/logging.h b/cpp/src/arrow/telemetry/logging.h
index 2ac0d8ea6b08b..30f91215e8762 100644
--- a/cpp/src/arrow/telemetry/logging.h
+++ b/cpp/src/arrow/telemetry/logging.h
@@ -97,7 +97,9 @@ struct EventId {
 struct LogDescriptor {
   LogLevel severity = kDefaultSeverity;
 
-  std::optional<std::string_view> body = std::nullopt;
+  std::chrono::system_clock::time_point timestamp = std::chrono::system_clock::now();
+
+  std::string_view body = "";
 
   EventId event_id = EventId::Invalid();
 
@@ -190,5 +192,22 @@ class ARROW_EXPORT GlobalLogger {
   static std::unique_ptr<Logger> logger_;
 };
 
+class ARROW_EXPORT LogMessage {
+ public:
+  explicit LogMessage(LogLevel, Logger*);
+
+  std::ostream& Stream();
+
+  template <typename T>
+  LogMessage& operator<<(const T& t) {
+    Stream() << t;
+    return *this;
+  }
+
+ private:
+  class Impl;
+  std::shared_ptr<Impl> impl_;
+};
+
 }  // namespace telemetry
 }  // namespace arrow
diff --git a/cpp/src/arrow/telemetry/telemetry_test.cc b/cpp/src/arrow/telemetry/telemetry_test.cc
index 362cc4dd74b9d..5e494aa268afd 100644
--- a/cpp/src/arrow/telemetry/telemetry_test.cc
+++ b/cpp/src/arrow/telemetry/telemetry_test.cc
@@ -50,6 +50,7 @@ class OtelEnvironment : public ::testing::Environment {
     auto provider_options = LoggerProviderOptions::Defaults();
     ASSERT_OK(GlobalLoggerProvider::Initialize(provider_options));
     auto logging_options = LoggingOptions::Defaults();
+    logging_options.severity_threshold = LogLevel::ARROW_TRACE;
     logging_options.flush_severity = LogLevel::ARROW_TRACE;
     ASSERT_OK_AND_ASSIGN(
         auto logger,
@@ -90,6 +91,8 @@ TEST_F(TestLogging, Basics) {
   Log(LogLevel::ARROW_WARNING, "baz bal",
       AttributeList{Attribute{"intAttr", 24}, Attribute{"boolAttr", true},
                     Attribute{"strAttr", std::string("ab") + "c"}});
+  LogMessage(LogLevel::ARROW_INFO, GlobalLogger::Get()) << "This is a "
+                                                        << "log message";
 }
 
 }  // namespace telemetry