feat: add flag to regex filter record types for SD targets

Addresses #3

cmd/ns1_exporter/ns1_exporter.go

@@ -128,6 +128,11 @@ var (
 		"A regular expression of zone(s) that the service discovery mechanism will provide targets for",
 	).Default("").Regexp()
 
+	flagNS1SDRecordTypeRegex = kingpin.Flag(
+		"ns1.sd-record-type",
+		"A regular expression of record types that the service discovery mechanism will provide targets for",
+	).Default("").Regexp()
+
 	flagRuntimeGOMAXPROCS = kingpin.Flag(
 		"runtime.gomaxprocs", "The target number of CPUs Go will run on (GOMAXPROCS)",
 	).Envar("GOMAXPROCS").Default("1").Int()
@@ -162,7 +167,7 @@ func Run() {
 		// EnableDDI:   *flagNS1EnableDDI,
 	})
 	exporterWorker := exporter.NewWorker(apiClient, *flagNS1ExporterEnableZoneQPS, *flagNS1ExporterEnableRecordQPS, *flagNS1ExporterZoneBlacklistRegex, *flagNS1ExporterZoneWhitelistRegex)
-	sdWorker := sd.NewWorker(apiClient, *flagNS1SDZoneBlacklistRegex, *flagNS1SDZoneWhitelistRegex)
+	sdWorker := sd.NewWorker(apiClient, *flagNS1SDZoneBlacklistRegex, *flagNS1SDZoneWhitelistRegex, *flagNS1SDRecordTypeRegex)
 
 	var g run.Group
 	{

pkg/servicediscovery/sd.go

@@ -67,20 +67,22 @@ type HTTPSDTarget struct {
 // Worker gets registered on a different handler for the `/sd` path and run via
 // the same HTTP server as the metrics exporter.
 type Worker struct {
-	ZoneBlacklist *regexp.Regexp
-	ZoneWhitelist *regexp.Regexp
+	ZoneBlacklist       *regexp.Regexp
+	ZoneWhitelist       *regexp.Regexp
+	RecordTypeWhitelist *regexp.Regexp
 
 	client      *api.Client
 	zoneCache   map[string]*ns1_internal.Zone
 	recordCache []*dns.Record
 	targetCache []*HTTPSDTarget
 }
 
-func NewWorker(client *api.Client, blacklist, whitelist *regexp.Regexp) *Worker {
+func NewWorker(client *api.Client, blacklist, whitelist, recordType *regexp.Regexp) *Worker {
 	worker := Worker{
-		client:        client,
-		ZoneBlacklist: blacklist,
-		ZoneWhitelist: whitelist,
+		client:              client,
+		ZoneBlacklist:       blacklist,
+		ZoneWhitelist:       whitelist,
+		RecordTypeWhitelist: recordType,
 	}
 
 	return &worker
@@ -241,10 +243,25 @@ func (w *Worker) RefreshRecordData() {
 	var records []*dns.Record
 
 	for zName, zData := range w.zoneCache {
-		for _, zRecord := range zData.Records {
-			record, _, err := w.client.Records.Get(zData.Zone, zRecord.Domain, zRecord.Type)
+		// if record type regex is provided, filter records
+		if w.RecordTypeWhitelist != nil && w.RecordTypeWhitelist.String() != "" {
+			var filteredRecords []*ns1_internal.ZoneRecord
+			for _, r := range zData.Records {
+				if !w.RecordTypeWhitelist.MatchString(r.Type) {
+					// if record type not in whitelist, log it and skip it
+					level.Debug(logger).Log("msg", "skipping record because it doesn't match whitelist regex", "worker", "http_sd", "record", r.Domain, "record_type_regex", w.RecordTypeWhitelist.String())
+					continue
+				}
+				filteredRecords = append(filteredRecords, r)
+			}
+
+			zData.Records = filteredRecords
+		}
+
+		for _, r := range zData.Records {
+			record, _, err := w.client.Records.Get(zData.Zone, r.Domain, r.Type)
 			if err != nil {
-				level.Error(logger).Log("msg", "Failed to get record data from NS1 API", "err", err.Error(), "worker", "http_sd", "zone_name", zName, "record_domain", zRecord.Domain, "record_type", zRecord.Type)
+				level.Error(logger).Log("msg", "Failed to get record data from NS1 API", "err", err.Error(), "worker", "http_sd", "zone_name", zName, "record_domain", r.Domain, "record_type", r.Type)
 				metrics.MetricExporterNS1APIFailures.Inc()
 				continue
 			}

pkg/servicediscovery/sd_test.go

@@ -1,4 +1,4 @@
-// Copyright 2023 TJ Hoplock
+// Copyright 2024 TJ Hoplock
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"regexp"
 	"testing"
 	"time"
 
@@ -218,7 +219,7 @@ func TestRecordAsPrometheusTarget(t *testing.T) {
 	mockClient.Endpoint, err = url.Parse(fmt.Sprintf("https://%s/v1/", mock.Address))
 	require.NoError(t, err)
 
-	worker := NewWorker(mockClient, nil, nil)
+	worker := NewWorker(mockClient, nil, nil, nil)
 
 	tests := map[string]struct {
 		recordCache []*dns.Record
@@ -254,17 +255,18 @@ func TestRefreshRecordData(t *testing.T) {
 	mockClient.Endpoint, err = url.Parse(fmt.Sprintf("https://%s/v1/", mock.Address))
 	require.NoError(t, err)
 
-	worker := NewWorker(mockClient, nil, nil)
-
 	tests := map[string]struct {
-		zoneCache map[string]*ns1_internal.Zone
-		want      []*dns.Record
+		zoneCache           map[string]*ns1_internal.Zone
+		recordTypeWhitelist *regexp.Regexp
+		want                []*dns.Record
 	}{
-		"empty_zone_cache": {zoneCache: map[string]*ns1_internal.Zone{}, want: nil},
-		"some_zone_cache":  {zoneCache: mockZoneCache, want: mockDnsRecordCache},
+		"empty_zone_cache":   {zoneCache: map[string]*ns1_internal.Zone{}, recordTypeWhitelist: nil, want: nil},
+		"some_zone_cache":    {zoneCache: mockZoneCache, recordTypeWhitelist: nil, want: mockDnsRecordCache},
+		"record_type_filter": {zoneCache: mockZoneCache, recordTypeWhitelist: regexp.MustCompile("SRV|AAAA"), want: mockDnsRecordCache[1:]},
 	}
 
 	for name, tc := range tests {
+		worker := NewWorker(mockClient, nil, nil, tc.recordTypeWhitelist)
 		worker.zoneCache = tc.zoneCache
 
 		t.Run(name, func(t *testing.T) {
@@ -293,7 +295,7 @@ func TestRefreshPrometheusTargetData(t *testing.T) {
 	mockClient.Endpoint, err = url.Parse(fmt.Sprintf("https://%s/v1/", mock.Address))
 	require.NoError(t, err)
 
-	worker := NewWorker(mockClient, nil, nil)
+	worker := NewWorker(mockClient, nil, nil, nil)
 
 	tests := map[string]struct {
 		recordCache []*dns.Record
@@ -329,7 +331,7 @@ func TestServeHTTP(t *testing.T) {
 	ts := httptest.NewServer(http.DefaultServeMux)
 	t.Cleanup(ts.Close)
 
-	worker := NewWorker(mockClient, nil, nil)
+	worker := NewWorker(mockClient, nil, nil, nil)
 	http.Handle("/sd", worker)
 	httpClient := http.Client{
 		Timeout: 30 * time.Second,