From d2cefbf22495b3c62956a455c6bbc63daf40c96b Mon Sep 17 00:00:00 2001
From: Jan Grewe <jan@faked.org>
Date: Mon, 30 May 2022 21:29:57 +0200
Subject: [PATCH] Allow Prometheus component to export internal components
 (#3508)

Co-authored-by: Jan Grewe <jan.grewe@flixbus.com>
---
 esphome/components/prometheus/__init__.py          |  8 +++++++-
 .../components/prometheus/prometheus_handler.cpp   | 14 +++++++-------
 esphome/components/prometheus/prometheus_handler.h |  8 ++++++++
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/esphome/components/prometheus/__init__.py b/esphome/components/prometheus/__init__.py
index 45345f06e8..e7c0459251 100644
--- a/esphome/components/prometheus/__init__.py
+++ b/esphome/components/prometheus/__init__.py
@@ -1,6 +1,9 @@
 import esphome.config_validation as cv
 import esphome.codegen as cg
-from esphome.const import CONF_ID
+from esphome.const import (
+    CONF_ID,
+    CONF_INCLUDE_INTERNAL,
+)
 from esphome.components.web_server_base import CONF_WEB_SERVER_BASE_ID
 from esphome.components import web_server_base
 
@@ -15,6 +18,7 @@ CONFIG_SCHEMA = cv.Schema(
         cv.GenerateID(CONF_WEB_SERVER_BASE_ID): cv.use_id(
             web_server_base.WebServerBase
         ),
+        cv.Optional(CONF_INCLUDE_INTERNAL, default=False): cv.boolean,
     },
     cv.only_with_arduino,
 ).extend(cv.COMPONENT_SCHEMA)
@@ -27,3 +31,5 @@ async def to_code(config):
 
     var = cg.new_Pvariable(config[CONF_ID], paren)
     await cg.register_component(var, config)
+
+    cg.add(var.set_include_internal(config[CONF_INCLUDE_INTERNAL]))
diff --git a/esphome/components/prometheus/prometheus_handler.cpp b/esphome/components/prometheus/prometheus_handler.cpp
index e4dd6b9043..a52347ba57 100644
--- a/esphome/components/prometheus/prometheus_handler.cpp
+++ b/esphome/components/prometheus/prometheus_handler.cpp
@@ -61,7 +61,7 @@ void PrometheusHandler::sensor_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_sensor_failed GAUGE\n"));
 }
 void PrometheusHandler::sensor_row_(AsyncResponseStream *stream, sensor::Sensor *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   if (!std::isnan(obj->state)) {
     // We have a valid value, output this value
@@ -98,7 +98,7 @@ void PrometheusHandler::binary_sensor_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_binary_sensor_failed GAUGE\n"));
 }
 void PrometheusHandler::binary_sensor_row_(AsyncResponseStream *stream, binary_sensor::BinarySensor *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   if (obj->has_state()) {
     // We have a valid value, output this value
@@ -134,7 +134,7 @@ void PrometheusHandler::fan_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_fan_oscillation GAUGE\n"));
 }
 void PrometheusHandler::fan_row_(AsyncResponseStream *stream, fan::Fan *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   stream->print(F("esphome_fan_failed{id=\""));
   stream->print(obj->get_object_id().c_str());
@@ -179,7 +179,7 @@ void PrometheusHandler::light_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_light_effect_active GAUGE\n"));
 }
 void PrometheusHandler::light_row_(AsyncResponseStream *stream, light::LightState *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   // State
   stream->print(F("esphome_light_state{id=\""));
@@ -255,7 +255,7 @@ void PrometheusHandler::cover_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_cover_failed GAUGE\n"));
 }
 void PrometheusHandler::cover_row_(AsyncResponseStream *stream, cover::Cover *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   if (!std::isnan(obj->position)) {
     // We have a valid value, output this value
@@ -298,7 +298,7 @@ void PrometheusHandler::switch_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_switch_failed GAUGE\n"));
 }
 void PrometheusHandler::switch_row_(AsyncResponseStream *stream, switch_::Switch *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   stream->print(F("esphome_switch_failed{id=\""));
   stream->print(obj->get_object_id().c_str());
@@ -322,7 +322,7 @@ void PrometheusHandler::lock_type_(AsyncResponseStream *stream) {
   stream->print(F("#TYPE esphome_lock_failed GAUGE\n"));
 }
 void PrometheusHandler::lock_row_(AsyncResponseStream *stream, lock::Lock *obj) {
-  if (obj->is_internal())
+  if (obj->is_internal() && !this->include_internal_)
     return;
   stream->print(F("esphome_lock_failed{id=\""));
   stream->print(obj->get_object_id().c_str());
diff --git a/esphome/components/prometheus/prometheus_handler.h b/esphome/components/prometheus/prometheus_handler.h
index 5c8d51c60f..b378e46ea3 100644
--- a/esphome/components/prometheus/prometheus_handler.h
+++ b/esphome/components/prometheus/prometheus_handler.h
@@ -13,6 +13,13 @@ class PrometheusHandler : public AsyncWebHandler, public Component {
  public:
   PrometheusHandler(web_server_base::WebServerBase *base) : base_(base) {}
 
+  /** Determine whether internal components should be exported as metrics.
+   * Defaults to false.
+   *
+   * @param include_internal Whether internal components should be exported.
+   */
+  void set_include_internal(bool include_internal) { include_internal_ = include_internal; }
+
   bool canHandle(AsyncWebServerRequest *request) override {
     if (request->method() == HTTP_GET) {
       if (request->url() == "/metrics")
@@ -84,6 +91,7 @@ class PrometheusHandler : public AsyncWebHandler, public Component {
 #endif
 
   web_server_base::WebServerBase *base_;
+  bool include_internal_{false};
 };
 
 }  // namespace prometheus