98 lines
2 KiB
Go
98 lines
2 KiB
Go
package main
|
|
|
|
import (
|
|
"database/sql"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/prometheus/client_golang/prometheus/promhttp"
|
|
_ "github.com/lib/pq"
|
|
_ "github.com/mattn/go-sqlite3"
|
|
)
|
|
|
|
const (
|
|
port = ":3000"
|
|
)
|
|
|
|
var (
|
|
db *sql.DB
|
|
|
|
loginCount = prometheus.NewGauge(prometheus.GaugeOpts{
|
|
Name: "pocket_id_login_count",
|
|
Help: "Total number of sign-ins",
|
|
})
|
|
|
|
userCount = prometheus.NewGauge(prometheus.GaugeOpts{
|
|
Name: "pocket_id_user_count",
|
|
Help: "Total number of users",
|
|
})
|
|
)
|
|
|
|
func init() {
|
|
prometheus.MustRegister(loginCount)
|
|
prometheus.MustRegister(userCount)
|
|
}
|
|
|
|
func main() {
|
|
dbType := os.Getenv("DB_TYPE")
|
|
dbConnection := os.Getenv("DB_CONNECTION")
|
|
|
|
if dbType == "" || dbConnection == "" {
|
|
log.Fatal("DB_TYPE and DB_CONNECTION environment variables must be set")
|
|
}
|
|
|
|
var err error
|
|
db, err = sql.Open(dbType, dbConnection)
|
|
if err != nil {
|
|
log.Fatalf("Error opening database: %v", err)
|
|
}
|
|
defer db.Close()
|
|
|
|
err = db.Ping()
|
|
if err != nil {
|
|
log.Fatalf("Error connecting to the database: %v", err)
|
|
}
|
|
|
|
log.Printf("Connected to the %s database.", dbType)
|
|
|
|
http.HandleFunc("/metrics", metricsHandler)
|
|
http.Handle("/", promhttp.Handler())
|
|
|
|
log.Printf("Server running on http://localhost%s", port)
|
|
log.Fatal(http.ListenAndServe(port, nil))
|
|
}
|
|
|
|
func metricsHandler(w http.ResponseWriter, r *http.Request) {
|
|
updateMetrics()
|
|
promhttp.Handler().ServeHTTP(w, r)
|
|
}
|
|
|
|
func updateMetrics() {
|
|
var count int
|
|
|
|
query := "SELECT COUNT(*) FROM Audit_Logs WHERE event = 'SIGN_IN'"
|
|
if os.Getenv("DB_TYPE") == "postgres" {
|
|
query = "SELECT COUNT(*) FROM audit_logs WHERE event = 'SIGN_IN'"
|
|
}
|
|
|
|
err := db.QueryRow(query).Scan(&count)
|
|
if err != nil {
|
|
log.Printf("Error querying login count: %v", err)
|
|
} else {
|
|
loginCount.Set(float64(count))
|
|
}
|
|
|
|
query = "SELECT COUNT(*) FROM Users"
|
|
if os.Getenv("DB_TYPE") == "postgres" {
|
|
query = "SELECT COUNT(*) FROM users"
|
|
}
|
|
|
|
err = db.QueryRow(query).Scan(&count)
|
|
if err != nil {
|
|
log.Printf("Error querying user count: %v", err)
|
|
} else {
|
|
userCount.Set(float64(count))
|
|
}
|
|
}
|