digitalocean.go

package digitalocean

import (
	"fmt"
	"time"

	"github.com/grafana/agent/component"
	"github.com/grafana/agent/component/common/config"
	"github.com/grafana/agent/component/discovery"
	"github.com/grafana/agent/pkg/river/rivertypes"
	"github.com/prometheus/common/model"
	prom_discovery "github.com/prometheus/prometheus/discovery/digitalocean"
)

func init() {
	component.Register(component.Registration{
		Name:    "discovery.digitalocean",
		Args:    Arguments{},
		Exports: discovery.Exports{},

		Build: func(opts component.Options, args component.Arguments) (component.Component, error) {
			return New(opts, args.(Arguments))
		},
	})
}

type Arguments struct {
	RefreshInterval time.Duration `river:"refresh_interval,attr,optional"`
	Port            int           `river:"port,attr,optional"`

	BearerToken     rivertypes.Secret `river:"bearer_token,attr,optional"`
	BearerTokenFile string            `river:"bearer_token_file,attr,optional"`

	ProxyURL        config.URL `river:"proxy_url,attr,optional"`
	FollowRedirects bool       `river:"follow_redirects,attr,optional"`
	EnableHTTP2     bool       `river:"enable_http2,attr,optional"`
}

var DefaultArguments = Arguments{
	Port:            80,
	RefreshInterval: time.Minute,
}

func (a *Arguments) UnmarshalRiver(f func(interface{}) error) error {
	*a = DefaultArguments
	type arguments Arguments
	err := f((*arguments)(a))
	if err != nil {
		return err
	}
	return a.Validate()
}

// Validate validates the arguments. Specifically, it checks that a BearerToken or
// BearerTokenFile is specified, as the DigitalOcean API requires a Bearer Token for
// authentication.
func (a *Arguments) Validate() error {
	if (a.BearerToken == "" && a.BearerTokenFile == "") ||
		(len(a.BearerToken) > 0 && len(a.BearerTokenFile) > 0) {

		return fmt.Errorf("exactly one of bearer_token or bearer_token_file must be specified")
	}

	return nil
}

func (a *Arguments) Convert() *prom_discovery.SDConfig {
	httpClientConfig := config.DefaultHTTPClientConfig
	httpClientConfig.BearerToken = a.BearerToken
	httpClientConfig.BearerTokenFile = a.BearerTokenFile
	httpClientConfig.ProxyURL = a.ProxyURL
	httpClientConfig.FollowRedirects = a.FollowRedirects
	httpClientConfig.EnableHTTP2 = a.EnableHTTP2

	return &prom_discovery.SDConfig{
		RefreshInterval:  model.Duration(a.RefreshInterval),
		Port:             a.Port,
		HTTPClientConfig: *httpClientConfig.Convert(),
	}
}

func New(opts component.Options, args Arguments) (component.Component, error) {
	return discovery.New(opts, args, func(args component.Arguments) (discovery.Discoverer, error) {
		newArgs := args.(Arguments)
		return prom_discovery.NewDiscovery(newArgs.Convert(), opts.Logger)
	})
}