Fix the prometheus type badges

handlers/badges.go

@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -14,8 +15,10 @@ import (
 	"github.com/jinzhu/gorm"
 	"github.com/jinzhu/gorm/dialects/postgres"
 	"github.com/prometheus/client_golang/api"
+	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prometheus/common/model"
 	"github.com/sirupsen/logrus"
 	"github.com/wahyd4/badger/models"
 )
@@ -39,6 +42,7 @@ const (
 	ColorSchemeOrange         = ColorScheme("fill-opacity:1.00; fill:rgb(255, 87, 34);")
 	ColorSchemeBlue           = ColorScheme("fill-opacity:1.00; fill:rgb(32, 99, 155);")
 	BadgeTypePrometheus       = "PROMETHEUS"
+	BadgeTypePageView         = "PAGE_VIEW"
 	pageViewsLabel            = "Page Views"
 	textMarginX               = 6.0
 	textMarginY               = 14
@@ -130,10 +134,21 @@ func (handler *BadgesHandler) generateSVG(c *gin.Context, colorScheme ColorSchem
 		sentry.CaptureException(err)
 		return err
 	}
-
+	var metric *Metric
 	prometheusOptions := badge.OptionStruct.(*models.PrometheusOptions)
+	if badge.Type == BadgeTypePrometheus {
+		metric, err = handler.queryPrometheus(badge, prometheusOptions)
+	} else {
+		metric, err = handler.queryMetric(badge, internalRequest)
+	}
 
-	metric, _ := handler.queryMetric(badge, internalRequest, prometheusOptions)
+	if err != nil {
+		logrus.Errorf("failed to query metric for badge %s : %v", badge.ULID, err)
+		metric = &Metric{
+			Label: badge.Label,
+			Value: -1,
+		}
+	}
 
 	metricValue := metric.Value
 	label := metric.Label
@@ -239,10 +254,10 @@ func (handler *BadgesHandler) CreateBadge(c *gin.Context) {
 	}
 
 	badge := models.Badge{
-		Label:   pageViewsLabel,
-		ULID:    newID,
-		Type:    BadgeTypePrometheus,
-		Options: jsonB,
+		Label: pageViewsLabel,
+		ULID:  newID,
+		Type:  BadgeTypePageView,
+		// Options: jsonB,
 	}
 
 	if c.GetBool("loggedIn") {
@@ -259,7 +274,7 @@ func (handler *BadgesHandler) CreateBadge(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{"id": newID})
 }
 
-func (handler *BadgesHandler) queryMetric(badge *models.Badge, internalRequest bool, options *models.PrometheusOptions) (*Metric, error) {
+func (handler *BadgesHandler) queryMetric(badge *models.Badge, internalRequest bool) (*Metric, error) {
 	viewCount, err := handler.fetchViewCountFromDB(badge)
 	if err != nil {
 		return nil, err
@@ -306,3 +321,29 @@ func (handler *BadgesHandler) updateMetric(viewCount *models.BadgeViewCount) err
 	handler.DB.Save(&models.ViewRecord{ULID: viewCount.BadgeULID})
 	return nil
 }
+
+func (handler *BadgesHandler) queryPrometheus(badge *models.Badge, options *models.PrometheusOptions) (*Metric, error) {
+	v1api := v1.NewAPI(handler.PrometheusAPIClient)
+	ctx, cancel := context.WithTimeout(context.Background(), metricRequestTimeout)
+	defer cancel()
+
+	result, warnings, err := v1api.Query(ctx, options.Query, time.Now())
+	if err != nil {
+		sentry.CaptureException(err)
+		return nil, fmt.Errorf("error querying Prometheus: %w", err)
+	}
+	if len(warnings) > 0 {
+		logrus.Warnf("Warnings: %v\n", warnings)
+	}
+
+	var metric Metric
+	for _, v := range result.(model.Vector) {
+		metric = Metric{
+			Label: string(v.Metric[model.LabelName(options.ValueField)]),
+			Value: int(float64(v.Value)),
+			Time:  v.Timestamp.Time(),
+		}
+	}
+	handler.DB.Save(&models.ViewRecord{ULID: badge.ULID})
+	return &metric, nil
+}

www/pages/stats/_id.vue

@@ -12,7 +12,7 @@
           <line-chart :data="chartsData" :options="options" v-if="loaded" />
         </b-col>
       </b-row>
-      <b-row class="">
+      <b-row class="get-your-badge">
         <b-col lg="4" offset-lg="4" class="text-center generate-badge">
           <b-button
             pill
@@ -20,7 +20,8 @@
             class="generate-button"
             variant="outline-primary"
             href="/"
-            >Create another Badge
+            >
+            Get your page views badge
           </b-button>
         </b-col>
       </b-row>
@@ -115,5 +116,8 @@ export default {
   .stats {
     margin-top: 2rem;
   }
+  .get-your-badge {
+    margin-top: 2rem;
+  }
 }
 </style>