Add output format Prometheus

src/Formats/registerFormats.cpp

@@ -72,6 +72,7 @@ void registerOutputFormatMySQLWire(FormatFactory & factory);
 void registerOutputFormatMarkdown(FormatFactory & factory);
 void registerOutputFormatPostgreSQLWire(FormatFactory & factory);
 void registerOutputFormatCapnProto(FormatFactory & factory);
+void registerOutputFormatPrometheus(FormatFactory & factory);
 
 /// Input only formats.
 
@@ -181,6 +182,7 @@ void registerFormats()
     registerOutputFormatMarkdown(factory);
     registerOutputFormatPostgreSQLWire(factory);
     registerOutputFormatCapnProto(factory);
+    registerOutputFormatPrometheus(factory);
 
     registerInputFormatRegexp(factory);
     registerInputFormatJSONAsString(factory);

src/Processors/Formats/Impl/PrometheusTextOutputFormat.cpp

@@ -0,0 +1,104 @@
+#include <optional>
+#include <type_traits>
+#include <Processors/Formats/Impl/PrometheusTextOutputFormat.h>
+#include <Formats/FormatFactory.h>
+#include <Formats/registerWithNamesAndTypes.h>
+#include <IO/WriteHelpers.h>
+#include <DataTypes/DataTypeNullable.h>
+#include <DataTypes/Serializations/ISerialization.h>
+#include <base/defines.h>
+
+namespace DB
+{
+
+constexpr auto FORMAT_NAME = "Prometheus";
+
+static bool isDataTypeString(const DataTypePtr & type)
+{
+    return WhichDataType(type).isStringOrFixedString();
+}
+
+template <typename ResType, typename Pred>
+static void getColumnPos(const Block & header, const String & col_name, Pred pred, ResType & res)
+{
+    static_assert(std::is_same_v<ResType, size_t> || std::is_same_v<ResType, std::optional<size_t>>, "Illegal ResType");
+
+    constexpr bool is_optional = std::is_same_v<ResType, std::optional<size_t>>;
+
+    if (header.has(col_name))
+    {
+        res = header.getPositionByName(col_name);
+        const auto & col = header.getByName(col_name);
+        if (!pred(is_optional ? removeNullable(col.type) : col.type))
+        {
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Illegal type '{}' of column '{}' for output format '{}'",
+                col.type->getName(), col_name, FORMAT_NAME);
+        }
+    }
+    else
+    {
+        if constexpr (is_optional)
+            res = std::nullopt;
+        else
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Column '{}' is required for output format '{}'", col_name, FORMAT_NAME);
+    }
+}
+
+PrometheusTextOutputFormat::PrometheusTextOutputFormat(
+    WriteBuffer & out_,
+    const Block & header_,
+    const RowOutputFormatParams & params_,
+    const FormatSettings & format_settings_)
+    : IRowOutputFormat(header_, out_, params_), format_settings(format_settings_)
+{
+    const Block & header = getPort(PortKind::Main).getHeader();
+
+    getColumnPos(header, "name", isDataTypeString, pos.name);
+    getColumnPos(header, "value", isNumber<DataTypePtr>, pos.value);
+
+    getColumnPos(header, "help", isDataTypeString,pos.help);
+    getColumnPos(header, "type", isDataTypeString, pos.type);
+}
+
+void PrometheusTextOutputFormat::write(const Columns & columns, size_t row_num)
+{
+    if (pos.help.has_value() && !columns[*pos.help]->isNullAt(row_num))
+    {
+        writeCString("# HELP ", out);
+        serializations[pos.name]->serializeText(*columns[pos.name], row_num, out, format_settings);
+        writeChar(' ', out);
+        serializations[*pos.help]->serializeText(*columns[*pos.help], row_num, out, format_settings);
+        writeChar('\n', out);
+    }
+
+    if (pos.type.has_value() && !columns[*pos.type]->isNullAt(row_num))
+    {
+        writeCString("# TYPE ", out);
+        serializations[pos.name]->serializeText(*columns[pos.name], row_num, out, format_settings);
+        writeChar(' ', out);
+        serializations[*pos.type]->serializeText(*columns[*pos.type], row_num, out, format_settings);
+        /// TODO(vdimir): Check if type is 'counter', 'gauge', 'histogram', 'summary', or 'untyped'
+        writeChar('\n', out);
+    }
+
+    serializations[pos.name]->serializeText(*columns[pos.name], row_num, out, format_settings);
+    writeChar(' ', out);
+    serializations[pos.value]->serializeText(*columns[pos.value], row_num, out, format_settings);
+
+    writeChar('\n', out);
+    writeChar('\n', out);
+}
+
+void registerOutputFormatPrometheus(FormatFactory & factory)
+{
+    factory.registerOutputFormat(FORMAT_NAME, [](
+        WriteBuffer & buf,
+        const Block & sample,
+        const RowOutputFormatParams & params,
+        const FormatSettings & settings)
+    {
+        return std::make_shared<PrometheusTextOutputFormat>(buf, sample, params, settings);
+    });
+}
+
+}

src/Processors/Formats/Impl/PrometheusTextOutputFormat.h

@@ -0,0 +1,51 @@
+#pragma once
+
+#include <Core/Block.h>
+#include <Formats/FormatSettings.h>
+#include <Processors/Formats/IRowOutputFormat.h>
+
+
+namespace DB
+{
+
+class WriteBuffer;
+
+/** A stream for outputting data in tsv format.
+  */
+class PrometheusTextOutputFormat : public IRowOutputFormat
+{
+public:
+    /** with_names - output in the first line a header with column names
+      * with_types - output the next line header with the names of the types
+      */
+    PrometheusTextOutputFormat(
+        WriteBuffer & out_,
+        const Block & header_,
+        const RowOutputFormatParams & params_,
+        const FormatSettings & format_settings_);
+
+    String getName() const override { return "PrometheusTextOutputFormat"; }
+
+    /// https://github.com/prometheus/docs/blob/86386ed25bc8a5309492483ec7d18d0914043162/content/docs/instrumenting/exposition_formats.md
+    String getContentType() const override { return "text/plain; version=0.0.4; charset=UTF-8"; }
+
+protected:
+    void write(const Columns & columns, size_t row_num) override;
+    void writeField(const IColumn &, const ISerialization &, size_t) override {}
+
+
+    struct ColumnPositions
+    {
+        size_t name;
+        size_t value;
+        std::optional<size_t> help;
+        std::optional<size_t> type;
+    };
+
+
+    ColumnPositions pos;
+
+    const FormatSettings format_settings;
+};
+
+}

tests/queries/0_stateless/02267_output_format_prometheus.reference

@@ -0,0 +1,15 @@
+# HELP metric0 info 0
+# TYPE metric0 counter
+metric0 0
+
+metric1 1
+
+# HELP metric2 info 2
+metric2 2
+
+# TYPE metric3 counter
+metric3 3
+
+# HELP metric4 info 4
+metric4 4
+

tests/queries/0_stateless/02267_output_format_prometheus.sql

@@ -0,0 +1,7 @@
+SELECT
+    'metric' || toString(number) as name,
+    number as value,
+    if(number % 2 == 0, 'info '  || toString(number), NULL) as help,
+    if(number % 3 == 0, 'counter', NULL) as type
+FROM numbers(5)
+FORMAT Prometheus