JsonSink.h

#pragma once

#include "quill/bundled/fmt/base.h"
#include "quill/core/Attributes.h"
#include "quill/core/Filesystem.h"
#include "quill/core/LogLevel.h"
#include "quill/core/MacroMetadata.h"
#include "quill/sinks/FileSink.h"
#include "quill/sinks/StreamSink.h"

#include "quill/bundled/fmt/format.h"

#include <cstdint>
#include <cstring>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

QUILL_BEGIN_NAMESPACE

namespace detail
{
template <typename TBase>
class JsonSink : public TBase
{
public:
  using base_type = TBase;

  using base_type::base_type;

  ~JsonSink() override = default;

  QUILL_ATTRIBUTE_HOT void write_log(MacroMetadata const* log_metadata, uint64_t log_timestamp,
                                     std::string_view thread_id, std::string_view thread_name,
                                     std::string const& process_id, std::string_view logger_name,
                                     LogLevel log_level, std::string_view log_level_description,
                                     std::string_view log_level_short_code,
                                     std::vector<std::pair<std::string, std::string>> const* named_args,
                                     std::string_view log_message, std::string_view log_statement) override
  {
    char const* message_format = log_metadata->message_format();

    if (strchr(log_metadata->message_format(), '\n') != nullptr)
    {
      // The format string contains at least one new line and that would break the json message, it needs to be removed
      _format = log_metadata->message_format();

      for (size_t pos = 0; (pos = _format.find('\n', pos)) != std::string::npos; pos++)
      {
        _format.replace(pos, 1, " ");
      }

      // we do not want newlines in the json message, use the modified message_format
      message_format = _format.data();
    }

    _json_message.clear();

    generate_json_message(log_metadata, log_timestamp, thread_id, thread_name, process_id,
                          logger_name, log_level, log_level_description, log_level_short_code,
                          named_args, log_message, log_statement, message_format);

    _json_message.append(std::string_view{"}\n"});

    StreamSink::write_log(log_metadata, log_timestamp, thread_id, thread_name, process_id, logger_name, log_level,
                          log_level_description, log_level_short_code, named_args, std::string_view{},
                          std::string_view{_json_message.data(), _json_message.size()});
  }

  QUILL_ATTRIBUTE_HOT virtual void generate_json_message(
    MacroMetadata const* log_metadata, uint64_t log_timestamp, std::string_view thread_id,
    std::string_view , std::string const& ,
    std::string_view logger_name, LogLevel ,
    std::string_view log_level_description, std::string_view ,
    std::vector<std::pair<std::string, std::string>> const* named_args,
    std::string_view , std::string_view , char const* message_format)
  {
    _json_message.append(fmtquill::format(
      R"({{"timestamp":"{}","file_name":"{}","line":"{}","thread_id":"{}","logger":"{}","log_level":"{}","message":"{}")",
      std::to_string(log_timestamp), log_metadata->file_name(), log_metadata->line(), thread_id,
      logger_name, log_level_description, message_format));

    // Add args as key-values
    if (named_args)
    {
      for (auto const& [key, value] : *named_args)
      {
        _json_message.append(std::string_view{",\""});
        _json_message.append(key);
        _json_message.append(std::string_view{"\":\""});
        _json_message.append(value);
        _json_message.append(std::string_view{"\""});
      }
    }
  }

protected:
  fmtquill::memory_buffer _json_message;
  std::string _format;
};
} // namespace detail

class JsonFileSink : public detail::JsonSink<FileSink>
{
public:
  JsonFileSink(fs::path const& filename, FileSinkConfig const& config,
               FileEventNotifier file_event_notifier = FileEventNotifier{}, bool do_fopen = true)
    : detail::JsonSink<FileSink>(filename, static_cast<FileSinkConfig const&>(config),
                                 std::move(file_event_notifier), do_fopen)
  {
  }

  ~JsonFileSink() override = default;
};

class JsonConsoleSink : public detail::JsonSink<StreamSink>
{
public:
  JsonConsoleSink() : detail::JsonSink<StreamSink>("stdout", nullptr) {}
  ~JsonConsoleSink() override = default;
};

QUILL_END_NAMESPACE