yaml.go

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

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