Add env vars and Prometheus metrics

main.go

@@ -2,67 +2,119 @@ package main
 
 import (
 	"context"
-	"log"
-	"os"
-	"regexp"
-	"time"
-
-	"net/http"
-
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
+	"log"
+	"net/http"
+	"os"
+	"regexp"
+	"strings"
+	"time"
 )
 
-type ShortenerAttr struct {
+type shortenerAttrType struct {
 	Slug string
-	URL string
+	URL  string
 }
 
-type Elem struct {
-	Shortener ShortenerAttr
+type inventoryItemType struct {
+	Shortener shortenerAttrType
 }
 
-var reValid = regexp.MustCompile(os.Getenv("GOREDIRECT_REGEX"))
-var mongoUri string = os.Getenv("MONGO_URI")
-var databaseName string = os.Getenv("GOREDIRECT_DATABASE")
-var collectionName string = os.Getenv("GOREDIRECT_COLLECTION")
-var redirectFound string = os.Getenv("GOREDIRECT_NOT_FOUND")
+var reValid = regexp.MustCompile("^[a-zA-Z0-9]{4,6}$")
+var mongoURI string = "mongodb://127.0.0.1:27017/default?replicaSet=rs0"
+var collectionName string = "inventory"
+var redirectFound string = os.Getenv("GOREDIRECT_FOUND")
 var redirectNotFound string = os.Getenv("GOREDIRECT_NOT_FOUND")
 
-func wrapper(coll *mongo.Collection) func (w http.ResponseWriter, r *http.Request) {
+func wrapper(coll *mongo.Collection) func(w http.ResponseWriter, r *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
+		counterQueries.Inc()
 		slug := r.URL.Path[1:]
 		match := reValid.MatchString(slug)
-		if (!match) {
-			w.WriteHeader(400)
+		if !match {
+			counterInvalidSlug.Inc()
+			http.Error(w, "Invalid slug", 400)
+			return
 		}
 
-		var elem Elem
-		if err := coll.FindOne(context.TODO(), bson.M{"shortener.slug": slug}).Decode(&elem); err != nil {
-			log.Fatal(err)
+		var doc inventoryItemType
+		err := coll.FindOne(context.TODO(), bson.M{"shortener.slug": slug}).Decode(&doc)
+		if err != nil {
+			counterNotFound.Inc()
+			if redirectNotFound == "" {
+				http.NotFound(w, r)
+				return
+			}
+
+			http.Redirect(w, r, strings.Replace(redirectNotFound, "%s", slug, 1), 302)
+			return
 		}
-		log.Printf("Redirecting %s to %s\n", slug, elem.Shortener.URL)
-		http.Redirect(w, r, elem.Shortener.URL, 302)
+
+		counterFound.Inc()
+		var u = doc.Shortener.URL
+		http.Redirect(w, r, u, 302)
+		log.Printf("Redirecting %s to %s\n", slug, doc.Shortener.URL)
 	}
 }
 
+var (
+	counterQueries = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "goredirect_queries",
+		Help: "The total number of queries.",
+	})
+	counterInvalidSlug = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "goredirect_invalid_slug",
+		Help: "The total number of queries that did not match regex.",
+	})
+	counterNotFound = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "goredirect_not_found",
+		Help: "The total number of queries with unknown slug.",
+	})
+	counterFound = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "goredirect_found",
+		Help: "The total number of queries that resulted in successful redirect.",
+	})
+)
+
 func main() {
-	client, err := mongo.NewClient(options.Client().ApplyURI(mongoUri))
+	if val := os.Getenv("GOREDIRECT_REGEX"); val != "" {
+		reValid = regexp.MustCompile(val)
+	}
+
+	if val := os.Getenv("MONGO_URI"); val != "" {
+		mongoURI = val
+	}
+
+	if val := os.Getenv("GOREDIRECT_COLLECTION"); val != "" {
+		collectionName = val
+	}
+
+	cs, err := connstring.ParseAndValidate(mongoURI)
+	client, err := mongo.NewClient(options.Client().ApplyURI(mongoURI))
 	if err != nil {
 		log.Fatal(err)
 	}
-	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	err = client.Connect(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
-	coll := client.Database(databaseName).Collection(collectionName)
+	cancel()
+	coll := client.Database(cs.Database).Collection(collectionName)
 	defer client.Disconnect(ctx)
+	http.Handle("/metrics", promhttp.Handler())
 	http.HandleFunc("/", wrapper(coll))
+
 	log.Printf("Starting HTTP server\n")
 	err2 := http.ListenAndServe(":80", nil)
 	if err2 != nil {
 		log.Fatal("ListenAndServe: ", err2)
 	}
+
 }