main.go

package main

import (
	"bytes"
	"context"
	"database/sql"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"runtime/debug"
	"strings"
	"time"

	klogrus "github.com/go-kit/kit/log/logrus"
	"github.com/kardianos/service"
	"github.com/lawyzheng/daemon"
	"github.com/lawyzheng/lyhook"
	go_ora "github.com/sijms/go-ora/v2"
	"github.com/sirupsen/logrus"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/prometheus/exporter-toolkit/web"
	"gopkg.in/alecthomas/kingpin.v2"
	//Required for debugging
	//_ "net/http/pprof"
)

var (
	// Version will be set at build time.
	Version            = "0.0.0.dev"
	listenAddress      = kingpin.Flag("web.listen-address", "Address to listen on for web interface and telemetry. (env: LISTEN_ADDRESS)").Default(getEnv("LISTEN_ADDRESS", ":9161")).String()
	metricPath         = kingpin.Flag("web.telemetry-path", "Path under which to expose metrics. (env: TELEMETRY_PATH)").Default(getEnv("TELEMETRY_PATH", "/metrics")).String()
	defaultFileMetrics = kingpin.Flag("default.metrics", "File with default metrics in a TOML file. (env: DEFAULT_METRICS)").Default(getEnv("DEFAULT_METRICS", filepath.Join(getCurrentPath(), "default-metrics.toml"))).String()
	customMetrics      = kingpin.Flag("custom.metrics", "File that may contain various custom metrics in a TOML file. (env: CUSTOM_METRICS)").Default(getEnv("CUSTOM_METRICS", "")).String()
	queryTimeout       = kingpin.Flag("query.timeout", "Query timeout (in seconds). (env: QUERY_TIMEOUT)").Default(getEnv("QUERY_TIMEOUT", "5")).String()
	maxIdleConns       = kingpin.Flag("database.maxIdleConns", "Number of maximum idle connections in the connection pool. (env: DATABASE_MAXIDLECONNS)").Default(getEnv("DATABASE_MAXIDLECONNS", "0")).Int()
	maxOpenConns       = kingpin.Flag("database.maxOpenConns", "Number of maximum open connections in the connection pool. (env: DATABASE_MAXOPENCONNS)").Default(getEnv("DATABASE_MAXOPENCONNS", "10")).Int()
	username           = kingpin.Flag("database.username", "Database username (env: DATABASE_USERNAME)").Default(getEnv("DATABASE_USERNAME", "")).String()
	password           = kingpin.Flag("database.password", "Database password (env: DATABASE_PASSWORD)").Default(getEnv("DATABASE_PASSWORD", "")).String()
	hostname           = kingpin.Flag("database.hostname", "Database hostname (env: DATABASE_HOSTNAME)").Default(getEnv("DATABASE_HOSTNAME", "")).String()
	port               = kingpin.Flag("database.port", "Database port (env: DATABASE_PORT)").Default(getEnv("DATABASE_PORT", "1521")).Int()
	dbservice          = kingpin.Flag("database.service", "Database service (env: DATABASE_SERVICE)").Default(getEnv("DATABASE_SERVICE", "")).String()
	dsn                = kingpin.Flag("datasource", "Database connect string (env: DATA_SOURCE_NAME)").Default(getEnv("DATA_SOURCE_NAME", "")).String()

	daemonMode     = kingpin.Flag("daemon", "Run as a daemon").Short('d').Bool()
	silence        = kingpin.Flag("silence", "Run in silence mode, only logging to file").Short('s').Bool()
	scrapeInterval = kingpin.Flag("scrape.interval", "Interval between each scrape. Default is to scrape on collect requests.").Default("0s").Duration()
	tlsconfigFile  = kingpin.Flag("web.config", "Path to config yaml file that can enable TLS or authentication.").Default(filepath.Join(getCurrentPath(), "default-tlsconfig.yaml")).String()
	serviceName    = kingpin.Flag("service.name", "Name to display in service. Default is [OracleExporter]").Default("OracleExporter").String()
)

var (
	panicWriter io.Writer = os.Stderr
)

// Metric name parts.
const (
	namespace = "oracledb"
	exporter  = "exporter"
)

func panicLogger() {
	if r := recover(); r != nil {
		buf := new(bytes.Buffer)
		buf.WriteString(fmt.Sprintf("[PANIC] %s panic: %v \n", time.Now().Format(time.RFC3339), r))
		buf.Write(debug.Stack())
		buf.WriteString("\n\n")
		panicWriter.Write(buf.Bytes())
	}
}

// getEnv returns the value of an environment variable, or returns the provided fallback value
func getEnv(key, fallback string) string {
	if value, ok := os.LookupEnv(key); ok {
		return value
	}
	return fallback
}

func maskDsn(dsn string) string {
	parts := strings.Split(dsn, "@")
	if len(parts) > 1 {
		maskedUrl := "***@" + parts[1]
		return maskedUrl
	}
	return dsn
}

func connect(dsn string, logger *logrus.Logger) (*sql.DB, error) {
	logger.Debug("Launching connection: ", maskDsn(dsn))
	db, err := sql.Open("oracle", dsn)
	if err != nil {
		logger.Error("Error while connecting to", dsn)
		return nil, err
	}
	logger.Debug("set max idle connections to ", *maxIdleConns)
	db.SetMaxIdleConns(*maxIdleConns)
	logger.Debug("set max open connections to ", *maxOpenConns)
	db.SetMaxOpenConns(*maxOpenConns)
	logger.Debug("Successfully connected to: ", maskDsn(dsn))
	return db, nil
}

func ping(ctx context.Context, db *sql.DB) error {
	newCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	return db.PingContext(newCtx)
}

func getCurrentPath() string {
	if ex, err := os.Executable(); err == nil {
		return filepath.Dir(ex)
	}
	return "."
}

func getLogger(config *promlog.Config) *logrus.Logger {
	var formatter logrus.Formatter
	if config.Format != nil && config.Format.String() == "json" {
		formatter = &logrus.JSONFormatter{}
	} else {
		formatter = &logrus.TextFormatter{DisableColors: true}
	}
	logrus.SetFormatter(formatter)

	f, _ := lyhook.NewRotateFile(filepath.Join(getCurrentPath(), "logs", "oracle_exporter.log"))
	lyhook.SetDefaultWriter(f)
	lyhook.SetFormatter(formatter)
	lyhook.SetCaller(lyhook.NewDefaultCaller().SetSkip(lyhook.DefaultCallerSkip + 1))

	plogger, _ := lyhook.NewRotateFile(filepath.Join(getCurrentPath(), "logs", "panic.log"))

	lv := logrus.InfoLevel
	if config.Level != nil {
		l, err := logrus.ParseLevel(config.Level.String())
		if err == nil {
			lv = l
		}
	}
	logrus.SetLevel(lv)

	if *silence {
		logrus.SetOutput(io.Discard)
		panicWriter = plogger
	} else {
		panicWriter = io.MultiWriter(panicWriter, plogger)
	}

	return logrus.StandardLogger()
}

func main() {
	promlogConfig := &promlog.Config{}
	flag.AddFlags(kingpin.CommandLine, promlogConfig)

	kingpin.Version("oracledb_exporter " + Version)
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	logger := getLogger(promlogConfig)
	s := NewServer(logger)

	var (
		svcConfig = &service.Config{
			Name:             *serviceName,
			DisplayName:      *serviceName,
			Description:      "A Prometheus exporter for Oracle modeled",
			WorkingDirectory: getCurrentPath(),
		}
	)

	if *daemonMode {
		envs := make(map[string]string)
		if s, ok := os.LookupEnv("LISTEN_ADDRESS"); ok {
			envs["LISTEN_ADDRESS"] = s
		}
		if s, ok := os.LookupEnv("TELEMETRY_PATH"); ok {
			envs["TELEMETRY_PATH"] = s
		}
		if s, ok := os.LookupEnv("DEFAULT_METRICS"); ok {
			envs["DEFAULT_METRICS"] = s
		}
		if s, ok := os.LookupEnv("CUSTOM_METRICS"); ok {
			envs["CUSTOM_METRICS"] = s
		}
		if s, ok := os.LookupEnv("QUERY_TIMEOUT"); ok {
			envs["QUERY_TIMEOUT"] = s
		}
		if s, ok := os.LookupEnv("DATABASE_MAXIDLECONNS"); ok {
			envs["DATABASE_MAXIDLECONNS"] = s
		}
		if s, ok := os.LookupEnv("DATABASE_MAXOPENCONNS"); ok {
			envs["DATABASE_MAXOPENCONNS"] = s
		}
		if s, ok := os.LookupEnv("DATA_SOURCE_NAME"); ok {
			envs["DATA_SOURCE_NAME"] = s
		}

		var args []string
		for _, arg := range os.Args[1:] {
			if strings.EqualFold(arg, "--daemon") || strings.EqualFold(arg, "-d") {
				continue
			}
			args = append(args, arg)
		}

		if !(*silence) {
			args = append(args, "-s")
		}

		svcConfig.EnvVars = envs
		svcConfig.Arguments = args

		ctl, err := daemon.NewController(svcConfig, s)
		if err != nil {
			logger.Error("new daemon error: ", err)
			return
		}

		ctl.Uninstall()
		if err := ctl.Install(); err != nil {
			logger.Error("install service error: ", err)
			return
		}

		if err := ctl.Start(); err != nil {
			logger.Error("start service error: ", err)
			return
		}

	} else {
		ctl, err := daemon.NewController(svcConfig, s)
		if err != nil {
			logger.Error("new daemon error: ", err)
			return
		}

		if err := ctl.Run(); err != nil {
			logger.Error("run application error: ", err)
			return
		}
	}
}

func NewServer(logger *logrus.Logger) *Server {
	return &Server{logger: logger}
}

type Server struct {
	logger *logrus.Logger
}

func (s *Server) Run() error {
	s.logger.Info("Starting oracledb_exporter " + Version)
	// dsn := os.Getenv("DATA_SOURCE_NAME")

	// Load default and custom metrics
	reloadMetrics(s.logger)

	var connstring string
	if strings.TrimSpace(*dsn) != "" {
		connstring = *dsn
	} else {
		connstring = go_ora.BuildUrl(*hostname, *port, *dbservice, *username, *password, nil)
	}

	exporter := NewExporter(connstring, s.logger)
	prometheus.MustRegister(exporter)
	go exporter.runScheduledScrapes()

	// See more info on https://github.com/prometheus/client_golang/blob/master/prometheus/promhttp/http.go#L269
	opts := promhttp.HandlerOpts{

		ErrorHandling: promhttp.ContinueOnError,
	}
	http.Handle(*metricPath, promhttp.HandlerFor(prometheus.DefaultGatherer, opts))

	http.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
		w.Write([]byte("<html><head><title>Oracle DB Exporter " + Version + "</title></head><body><h1>Oracle DB Exporter " + Version + "</h1><p><a href='" + *metricPath + "'>Metrics</a></p></body></html>"))
	})

	s.logger.Info("Listening on", *listenAddress)
	server := &http.Server{Addr: *listenAddress}
	return web.ListenAndServe(server, *tlsconfigFile, klogrus.NewLogger(s.logger, klogrus.WithLevel(s.logger.GetLevel())))
}

func (s *Server) HandleError(e error) {
	s.logger.Error("err: ", e)
	os.Exit(1)
}