yaml.go

Documentation: github.com/cybertec-postgresql/pgwatch/v5/internal/metrics

     1  package metrics
     2  
     3  import (
     4  	"context"
     5  	_ "embed"
     6  	"maps"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  	"sync"
    11  
    12  	"gopkg.in/yaml.v3"
    13  )
    14  
    15  func NewYAMLMetricReaderWriter(ctx context.Context, path string) (ReaderWriter, error) {
    16  	if path == "" {
    17  		return NewDefaultMetricReader(ctx)
    18  	}
    19  	return &fileMetricReader{
    20  		ctx:  ctx,
    21  		path: path,
    22  	}, nil
    23  }
    24  
    25  type fileMetricReader struct {
    26  	ctx  context.Context
    27  	path string
    28  	sync.Mutex
    29  }
    30  
    31  // WriteMetrics writes metrics to file with locking
    32  func (fmr *fileMetricReader) WriteMetrics(metricDefs *Metrics) error {
    33  	fmr.Lock()
    34  	defer fmr.Unlock()
    35  	return fmr.writeMetrics(metricDefs)
    36  }
    37  
    38  // writeMetrics writes metrics to file without locking (internal use only)
    39  func (fmr *fileMetricReader) writeMetrics(metricDefs *Metrics) error {
    40  	yamlData, _ := yaml.Marshal(metricDefs)
    41  	return os.WriteFile(fmr.path, yamlData, 0644)
    42  }
    43  
    44  //go:embed metrics.yaml
    45  var defaultMetricsYAML []byte
    46  
    47  // GetMetrics reads metrics from file or returns default metrics if path is empty with locking
    48  func (fmr *fileMetricReader) GetMetrics() (metrics *Metrics, err error) {
    49  	fmr.Lock()
    50  	defer fmr.Unlock()
    51  	return fmr.getMetrics()
    52  }
    53  
    54  // getMetrics reads metrics from file or returns default metrics if path is empty without locking (internal use only)
    55  func (fmr *fileMetricReader) getMetrics() (metrics *Metrics, err error) {
    56  	metrics = &Metrics{MetricDefs{}, PresetDefs{}}
    57  	if fmr.path == "" {
    58  		err = yaml.Unmarshal(defaultMetricsYAML, metrics)
    59  		return
    60  	}
    61  
    62  	fi, err := os.Stat(fmr.path)
    63  	if err != nil {
    64  		return nil, err
    65  	}
    66  	switch mode := fi.Mode(); {
    67  	case mode.IsDir():
    68  		var entries []os.DirEntry
    69  		if entries, err = os.ReadDir(fmr.path); err != nil {
    70  			return nil, err
    71  		}
    72  		for _, entry := range entries {
    73  			if ext := strings.ToLower(filepath.Ext(entry.Name())); ext != ".yaml" && ext != ".yml" {
    74  				continue
    75  			}
    76  			var m *Metrics
    77  			if m, err = fmr.loadMetricsFromFile(filepath.Join(fmr.path, entry.Name())); err != nil {
    78  				return nil, err
    79  			}
    80  			maps.Copy(metrics.PresetDefs, m.PresetDefs)
    81  			maps.Copy(metrics.MetricDefs, m.MetricDefs)
    82  		}
    83  	case mode.IsRegular():
    84  		metrics, err = fmr.loadMetricsFromFile(fmr.path)
    85  	}
    86  	return
    87  }
    88  
    89  // loadMetricsFromFile reads metrics from a single YAML file
    90  func (fmr *fileMetricReader) loadMetricsFromFile(metricsFilePath string) (metrics *Metrics, err error) {
    91  	var yamlFile []byte
    92  	if yamlFile, err = os.ReadFile(metricsFilePath); err != nil {
    93  		return
    94  	}
    95  	metrics = &Metrics{MetricDefs{}, PresetDefs{}}
    96  	err = yaml.Unmarshal(yamlFile, &metrics)
    97  	return
    98  }
    99  
   100  // DeleteMetric deletes a metric by name and writes the updated metrics back to file
   101  func (fmr *fileMetricReader) DeleteMetric(metricName string) error {
   102  	fmr.Lock()
   103  	defer fmr.Unlock()
   104  	metrics, err := fmr.getMetrics()
   105  	if err != nil {
   106  		return err
   107  	}
   108  	delete(metrics.MetricDefs, metricName)
   109  	return fmr.writeMetrics(metrics)
   110  }
   111  
   112  // UpdateMetric updates an existing metric or creates it if it doesn't exist, then writes the updated metrics back to file
   113  func (fmr *fileMetricReader) UpdateMetric(metricName string, metric Metric) error {
   114  	fmr.Lock()
   115  	defer fmr.Unlock()
   116  	metrics, err := fmr.getMetrics()
   117  	if err != nil {
   118  		return err
   119  	}
   120  	metrics.MetricDefs[metricName] = metric
   121  	return fmr.writeMetrics(metrics)
   122  }
   123  
   124  // CreateMetric creates a new metric if it doesn't already exist, then writes the updated metrics back to file
   125  func (fmr *fileMetricReader) CreateMetric(metricName string, metric Metric) error {
   126  	fmr.Lock()
   127  	defer fmr.Unlock()
   128  	metrics, err := fmr.getMetrics()
   129  	if err != nil {
   130  		return err
   131  	}
   132  	// Check if metric already exists
   133  	if _, exists := metrics.MetricDefs[metricName]; exists {
   134  		return ErrMetricExists
   135  	}
   136  	metrics.MetricDefs[metricName] = metric
   137  	return fmr.writeMetrics(metrics)
   138  }
   139  
   140  // DeletePreset deletes a preset by name and writes the updated metrics back to file
   141  func (fmr *fileMetricReader) DeletePreset(presetName string) error {
   142  	fmr.Lock()
   143  	defer fmr.Unlock()
   144  	metrics, err := fmr.getMetrics()
   145  	if err != nil {
   146  		return err
   147  	}
   148  	delete(metrics.PresetDefs, presetName)
   149  	return fmr.writeMetrics(metrics)
   150  }
   151  
   152  // UpdatePreset updates an existing preset or creates it if it doesn't exist, then writes the updated metrics back to file
   153  func (fmr *fileMetricReader) UpdatePreset(presetName string, preset Preset) error {
   154  	fmr.Lock()
   155  	defer fmr.Unlock()
   156  	metrics, err := fmr.getMetrics()
   157  	if err != nil {
   158  		return err
   159  	}
   160  	metrics.PresetDefs[presetName] = preset
   161  	return fmr.writeMetrics(metrics)
   162  }
   163  
   164  // CreatePreset creates a new preset if it doesn't already exist, then writes the updated metrics back to file
   165  func (fmr *fileMetricReader) CreatePreset(presetName string, preset Preset) error {
   166  	fmr.Lock()
   167  	defer fmr.Unlock()
   168  	metrics, err := fmr.getMetrics()
   169  	if err != nil {
   170  		return err
   171  	}
   172  	// Check if preset already exists
   173  	if _, exists := metrics.PresetDefs[presetName]; exists {
   174  		return ErrPresetExists
   175  	}
   176  	metrics.PresetDefs[presetName] = preset
   177  	return fmr.writeMetrics(metrics)
   178  }
   179