badger/main.go

package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/getsentry/sentry-go"
	"github.com/gin-gonic/gin"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	"github.com/prometheus/client_golang/api"
	"github.com/wahyd4/badger/handlers"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"

	sentrygin "github.com/getsentry/sentry-go/gin"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/sirupsen/logrus"
	"github.com/wahyd4/badger/config"
	"github.com/wahyd4/badger/models"
	gintrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/gin-gonic/gin"
)

type ColorScheme string

type queryParam struct {
	Color      ColorScheme
	Query      string
	ValueField string // nullable
	Label      string // when ValueField is null, then the program will use this field as the label
}

func main() {
	if err := sentry.Init(sentry.ClientOptions{
		Dsn: assertNotNullENV("SENTRY_DSN"),
	}); err != nil {
		log.Panicf("Sentry initialization failed: %v\n", err)
	}
	defer sentry.Flush(2 * time.Second)

	dbConfig := config.Config{
		Host:     getEnvWithDefault("DB_HOST", "new-postgres-postgresql.db.svc.cluster.local"),
		Port:     getEnvWithDefault("DB_PORT", "5432"),
		DBName:   getEnvWithDefault("DB_NAME", "badges_dev"),
		Username: getEnvWithDefault("DB_USERNAME", "someuser"),
		Password: getEnvWithDefault("DB_PASSWORD", "hahah"),
	}
	connStr := dbConfig.GetConnectionString()

	db, err := gorm.Open("postgres", connStr)
	if err != nil {
		sentry.CaptureException(err)
		panic(fmt.Sprintf("failed to connect database %s: %v", connStr, err))
	}
	defer db.Close()

	// Migrate the schema
	db.AutoMigrate(
		&models.Badge{},
		&models.Token{},
		&models.User{},
		&models.ViewRecord{},
		&models.FeatureToggle{},
		&models.BadgeViewCount{},
		&models.BadgeAPIValue{},
		&models.UserAPIToken{},
	)

	client, err := api.NewClient(api.Config{
		Address: "https://prometheus-api.junv.me",
	})
	if err != nil {
		sentry.CaptureException(err)
		logrus.Errorf("Error creating client: %v\n", err)
		os.Exit(1)
	}

	handler := &handlers.BadgesHandler{DB: db, PrometheusAPIClient: client}

	userHandler := &handlers.UsersHandler{DB: db}
	statsHandler := &handlers.StatsHandler{DB: db}
	toggleHandler := &handlers.FeatureTogglesHandler{DB: db}

	tracer.Start()
	defer tracer.Stop()

	r := gin.Default()

	r.Use(sentrygin.New(sentrygin.Options{
		Repanic: true,
	}))

	r.Use(gintrace.Middleware(getEnvWithDefault("DB_NAME", "badges_dev")))

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{"msg": "Hello!"})
	})

	r.GET("/svg/:name", handler.GetBadge)
	r.GET("/val/:name", handler.GetBadgeValue)
	r.PUT("/val/:name", handler.UpdateBadgeValue)

	r.GET("/badges/:name/:color", handler.GetBadgeWithColorScheme)

	api := r.Group("/api")
	api.Use(userHandler.AuthRequired)

	api.POST("/auth", userHandler.Authenticate)
	api.POST("/badges", handler.CreateBadge)
	api.GET("/badges", handler.MyBadges)
	api.GET("/toggles", toggleHandler.Toggles)

	statsAPI := r.Group("/api/stats")
	statsAPI.Use(handler.ValidateBadge)
	statsAPI.GET(":name", statsHandler.GetStats)

	go func() {
		http.Handle("/metrics", promhttp.Handler())
		if err := http.ListenAndServe(":9090", nil); err != nil {
			sentry.CaptureException(err)
			logrus.Error(err)
			os.Exit(1)
		}
	}()

	log.Fatal(r.Run())
}

func getEnvWithDefault(key, defaultValue string) string {
	val := os.Getenv(key)
	if val != "" {
		return val
	}
	return defaultValue
}

func assertNotNullENV(key string) string {
	val := os.Getenv(key)
	if val != "" {
		return val
	}
	panic(fmt.Sprintf("env %s must be set", key))
}