Added TSDB import with OpenMetrics and CSV file support.

Based on #5887 Thanks for your
work so far @dipack95, it helped a lot!

Changes on top of @dipack95:

* Addressed all reviews components
* Use subcommands for different formats
* Simplifed block creation, no need to be such complex for first iteration.
* Simpliefied and separate concerns. No need to have access to DB. Block
* writting is separated as well for ease of benchmarking and test. This will be also
needed by @JessicaGreben
* Added import support for different formats.
* Removed all tests - those had to be pulled over and adjusted ):


Signed-off-by: Bartlomiej Plotka <[email protected]>

Author: bwplotka

pkg/textparse/interface.go

@@ -16,6 +16,7 @@ package textparse
 import (
 	"mime"
 
+	"github.com/pkg/errors"
 	"github.com/prometheus/prometheus/pkg/exemplar"
 	"github.com/prometheus/prometheus/pkg/labels"
 )
@@ -55,8 +56,8 @@ type Parser interface {
 	// exemplar. It returns if an exemplar exists or not.
 	Exemplar(l *exemplar.Exemplar) bool
 
-	// Next advances the parser to the next sample. It returns false if no
-	// more samples were read or an error occurred.
+	// Next advances the parser to the next sample. It returns io.EOF if no
+	// more samples were read.
 	Next() (Entry, error)
 }
 
@@ -94,3 +95,45 @@ const (
 	MetricTypeStateset       = "stateset"
 	MetricTypeUnknown        = "unknown"
 )
+
+func (m *MetricType) ParseForOpenMetrics(mtyp string) error {
+	switch mtyp {
+	case "counter":
+		*m = MetricTypeCounter
+	case "gauge":
+		*m = MetricTypeGauge
+	case "histogram":
+		*m = MetricTypeHistogram
+	case "gaugehistogram":
+		*m = MetricTypeGaugeHistogram
+	case "summary":
+		*m = MetricTypeSummary
+	case "info":
+		*m = MetricTypeInfo
+	case "stateset":
+		*m = MetricTypeStateset
+	case "unknown":
+		*m = MetricTypeUnknown
+	default:
+		return errors.Errorf("invalid metric type %q", mtyp)
+	}
+	return nil
+}
+
+func (m *MetricType) ParseForProm(mtyp string) error {
+	switch mtyp {
+	case "counter":
+		*m = MetricTypeCounter
+	case "gauge":
+		*m = MetricTypeGauge
+	case "histogram":
+		*m = MetricTypeHistogram
+	case "summary":
+		*m = MetricTypeSummary
+	case "unknown":
+		*m = MetricTypeUnknown
+	default:
+		return errors.Errorf("invalid metric type %q", mtyp)
+	}
+	return nil
+}

pkg/textparse/openmetricsparse.go

@@ -259,25 +259,8 @@ func (p *OpenMetricsParser) Next() (Entry, error) {
 		}
 		switch t {
 		case tType:
-			switch s := yoloString(p.text); s {
-			case "counter":
-				p.mtype = MetricTypeCounter
-			case "gauge":
-				p.mtype = MetricTypeGauge
-			case "histogram":
-				p.mtype = MetricTypeHistogram
-			case "gaugehistogram":
-				p.mtype = MetricTypeGaugeHistogram
-			case "summary":
-				p.mtype = MetricTypeSummary
-			case "info":
-				p.mtype = MetricTypeInfo
-			case "stateset":
-				p.mtype = MetricTypeStateset
-			case "unknown":
-				p.mtype = MetricTypeUnknown
-			default:
-				return EntryInvalid, errors.Errorf("invalid metric type %q", s)
+			if err := p.mtype.ParseForOpenMetrics(yoloString(p.text)); err != nil {
+				return EntryInvalid, err
 			}
 		case tHelp:
 			if !utf8.Valid(p.text) {

pkg/textparse/promparse.go

@@ -289,19 +289,8 @@ func (p *PromParser) Next() (Entry, error) {
 		}
 		switch t {
 		case tType:
-			switch s := yoloString(p.text); s {
-			case "counter":
-				p.mtype = MetricTypeCounter
-			case "gauge":
-				p.mtype = MetricTypeGauge
-			case "histogram":
-				p.mtype = MetricTypeHistogram
-			case "summary":
-				p.mtype = MetricTypeSummary
-			case "untyped":
-				p.mtype = MetricTypeUnknown
-			default:
-				return EntryInvalid, errors.Errorf("invalid metric type %q", s)
+			if err := p.mtype.ParseForProm(yoloString(p.text)); err != nil {
+				return EntryInvalid, err
 			}
 		case tHelp:
 			if !utf8.Valid(p.text) {

tsdb/cmd/tsdb/main.go

@@ -34,11 +34,15 @@ import (
 	"github.com/go-kit/kit/log"
 	"github.com/pkg/errors"
 	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/pkg/textparse"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/prometheus/prometheus/tsdb/chunks"
 	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
 	"github.com/prometheus/prometheus/tsdb/importer"
+	"github.com/prometheus/prometheus/tsdb/importer/blocks"
+	"github.com/prometheus/prometheus/tsdb/importer/csv"
+	"github.com/prometheus/prometheus/tsdb/importer/openmetrics"
 	"gopkg.in/alecthomas/kingpin.v2"
 )
 
@@ -53,34 +57,43 @@ func execute() (err error) {
 	var (
 		defaultDBPath = filepath.Join("benchout", "storage")
 
-		cli                      = kingpin.New(filepath.Base(os.Args[0]), "CLI tool for tsdb")
-		benchCmd                 = cli.Command("bench", "run benchmarks")
-		benchWriteCmd            = benchCmd.Command("write", "run a write performance benchmark")
-		benchWriteOutPath        = benchWriteCmd.Flag("out", "set the output path").Default("benchout").String()
-		benchWriteNumMetrics     = benchWriteCmd.Flag("metrics", "number of metrics to read").Default("10000").Int()
-		benchSamplesFile         = benchWriteCmd.Arg("file", "input file with samples data, default is ("+filepath.Join("..", "..", "testdata", "20kseries.json")+")").Default(filepath.Join("..", "..", "testdata", "20kseries.json")).String()
-		listCmd                  = cli.Command("ls", "list db blocks")
-		listCmdHumanReadable     = listCmd.Flag("human-readable", "print human readable values").Short('h').Bool()
-		listPath                 = listCmd.Arg("db path", "database path (default is "+defaultDBPath+")").Default(defaultDBPath).String()
-		analyzeCmd               = cli.Command("analyze", "analyze churn, label pair cardinality.")
-		analyzePath              = analyzeCmd.Arg("db path", "database path (default is "+defaultDBPath+")").Default(defaultDBPath).String()
-		analyzeBlockID           = analyzeCmd.Arg("block id", "block to analyze (default is the last block)").String()
-		analyzeLimit             = analyzeCmd.Flag("limit", "how many items to show in each list").Default("20").Int()
-		dumpCmd                  = cli.Command("dump", "dump samples from a TSDB")
-		dumpPath                 = dumpCmd.Arg("db path", "database path (default is "+defaultDBPath+")").Default(defaultDBPath).String()
-		dumpMinTime              = dumpCmd.Flag("min-time", "minimum timestamp to dump").Default(strconv.FormatInt(math.MinInt64, 10)).Int64()
-		dumpMaxTime              = dumpCmd.Flag("max-time", "maximum timestamp to dump").Default(strconv.FormatInt(math.MaxInt64, 10)).Int64()
-		importCmd                = cli.Command("import", "import samples from file containing information formatted in the Open Metrics format. Please refer to the storage docs for more details.")
-		importFilePath           = importCmd.Arg("file path", "file to import samples from (must be in Open Metrics format)").Required().String()
-		importDbPath             = importCmd.Arg("db path", "database path").Required().String()
-		importMaxSamplesInMemory = importCmd.Flag("max-samples-in-mem", "maximum number of samples to process in a cycle").Default("10000").Int()
-		importMaxBlockChildren   = importCmd.Flag("max-block-children", "maximum number of children a block can have at a given time").Default("20").Int()
+		cli = kingpin.New(filepath.Base(os.Args[0]), "CLI tool for tsdb")
+
+		benchCmd             = cli.Command("bench", "run benchmarks")
+		benchWriteCmd        = benchCmd.Command("write", "run a write performance benchmark")
+		benchWriteOutPath    = benchWriteCmd.Flag("out", "set the output path").Default("benchout").String()
+		benchWriteNumMetrics = benchWriteCmd.Flag("metrics", "number of metrics to read").Default("10000").Int()
+		benchSamplesFile     = benchWriteCmd.Arg("file", "input file with samples data, default is ("+filepath.Join("..", "..", "testdata", "20kseries.json")+")").Default(filepath.Join("..", "..", "testdata", "20kseries.json")).String()
+
+		listCmd              = cli.Command("ls", "list db blocks")
+		listCmdHumanReadable = listCmd.Flag("human-readable", "print human readable values").Short('h').Bool()
+		listPath             = listCmd.Arg("db path", "database path (default is "+defaultDBPath+")").Default(defaultDBPath).String()
+
+		analyzeCmd     = cli.Command("analyze", "analyze churn, label pair cardinality.")
+		analyzePath    = analyzeCmd.Arg("db path", "database path (default is "+defaultDBPath+")").Default(defaultDBPath).String()
+		analyzeBlockID = analyzeCmd.Arg("block id", "block to analyze (default is the last block)").String()
+		analyzeLimit   = analyzeCmd.Flag("limit", "how many items to show in each list").Default("20").Int()
+
+		dumpCmd     = cli.Command("dump", "dump samples from a TSDB")
+		dumpPath    = dumpCmd.Arg("db path", "database path (default is "+defaultDBPath+")").Default(defaultDBPath).String()
+		dumpMinTime = dumpCmd.Flag("min-time", "minimum timestamp to dump").Default(strconv.FormatInt(math.MinInt64, 10)).Int64()
+		dumpMaxTime = dumpCmd.Flag("max-time", "maximum timestamp to dump").Default(strconv.FormatInt(math.MaxInt64, 10)).Int64()
+
+		importCmd       = cli.Command("import", "[Experimental] import samples from input and produce TSDB block. Please refer to the storage docs for more details.")
+		importDbPath    = importCmd.Flag("output", "output directory for generated block").Default(".").String()
+		importFilePath  = importCmd.Flag("input-file", "disables reading from input and using file to import samples from. If empty input is required").String()
+		importBlockSize = importCmd.Flag("block-size", "The maximum block size. The actual block timestamps will be aligned with Prometheus time ranges.").Default("2h").Hidden().Duration()
+
+		omImportCmd = importCmd.Command("openmetrics", "import samples from OpenMetrics input and produce TSDB block. Please refer to the storage docs for more details.")
+
+		csvImportCmd       = importCmd.Command("csv", "import samples from CSV input and produce TSDB block. Please refer to the storage docs for more details.")
+		csvImportDelimiter = csvImportCmd.Flag("delimiter", "CSV single character for fields delimiting").Default(",").String()
 	)
 
 	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
 	var merr tsdb_errors.MultiError
 
-	switch kingpin.MustParse(cli.Parse(os.Args[1:])) {
+	switch cmd := kingpin.MustParse(cli.Parse(os.Args[1:])); cmd {
 	case benchWriteCmd.FullCommand():
 		wb := &writeBenchmark{
 			outPath:     *benchWriteOutPath,
@@ -144,16 +157,39 @@ func execute() (err error) {
 			err = merr.Err()
 		}()
 		return dumpSamples(db, *dumpMinTime, *dumpMaxTime)
-	case importCmd.FullCommand():
-		f, err := os.Open(*importFilePath)
-		if err != nil {
-			return err
+	case omImportCmd.FullCommand(), csvImportCmd.FullCommand():
+		input := os.Stdin
+		if importFilePath != nil {
+			input, err = os.Open(*importFilePath)
+			if err != nil {
+				return err
+			}
+			defer func() {
+				merr.Add(err)
+				merr.Add(input.Close())
+				err = merr.Err()
+			}()
+		}
+
+		var p textparse.Parser
+		if cmd == omImportCmd.FullCommand() {
+			p = openmetrics.NewParser(input)
+		} else {
+			if len(*csvImportDelimiter) != 1 {
+				return errors.Errorf("wrong format of delimiter flag, expected single character, got %q", *csvImportDelimiter)
+			}
+
+			p = csv.NewParser(input, []rune(*csvImportDelimiter)[0])
 		}
-		return importer.ImportFromFile(f, *importDbPath, *importMaxSamplesInMemory, *importMaxBlockChildren, logger)
+		return importer.Import(logger, p, blocks.NewMultiWriter(logger, *importDbPath, durToMillis(*importBlockSize)))
 	}
 	return nil
 }
 
+func durToMillis(t time.Duration) int64 {
+	return int64(t.Seconds() * 1000)
+}
+
 type writeBenchmark struct {
 	outPath     string
 	samplesFile string

tsdb/importer/blocks/multi.go

@@ -0,0 +1,107 @@
+package blocks
+
+import (
+	"github.com/go-kit/kit/log"
+	"github.com/oklog/ulid"
+	"github.com/pkg/errors"
+	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/storage"
+	tsdb_errors "github.com/prometheus/prometheus/tsdb/errors"
+	"github.com/prometheus/prometheus/tsdb/index"
+)
+
+type errAppender struct{ err error }
+
+func (a errAppender) Add(l labels.Labels, t int64, v float64) (uint64, error) { return 0, a.err }
+func (a errAppender) AddFast(ref uint64, t int64, v float64) error            { return a.err }
+func (a errAppender) Commit() error                                           { return a.err }
+func (a errAppender) Rollback() error                                         { return a.err }
+
+func rangeForTimestamp(t int64, width int64) (maxt int64) {
+	return (t/width)*width + width
+}
+
+type MultiWriter struct {
+	blocks          map[index.Range]Writer
+	activeAppenders map[index.Range]storage.Appender
+
+	logger log.Logger
+	dir    string
+	// TODO(bwplotka): Allow more complex compaction levels.
+	sizeMillis int64
+}
+
+func NewMultiWriter(logger log.Logger, dir string, sizeMillis int64) *MultiWriter {
+	return &MultiWriter{
+		logger:          logger,
+		dir:             dir,
+		sizeMillis:      sizeMillis,
+		blocks:          map[index.Range]Writer{},
+		activeAppenders: map[index.Range]storage.Appender{},
+	}
+}
+
+// Appender is not thread-safe. Returned Appender is not thread-save as well.
+// TODO(bwplotka): Consider making it thread safe.
+func (w *MultiWriter) Appender() storage.Appender { return w }
+
+func (w *MultiWriter) getOrCreate(t int64) storage.Appender {
+	maxt := rangeForTimestamp(t, w.sizeMillis)
+	hash := index.Range{Start: maxt - w.sizeMillis, End: maxt}
+	if a, ok := w.activeAppenders[hash]; ok {
+		return a
+	}
+
+	nw, err := NewTSDBWriter(w.logger, w.dir)
+	if err != nil {
+		return errAppender{err: errors.Wrap(err, "new tsdb writer")}
+	}
+
+	w.blocks[hash] = nw
+	w.activeAppenders[hash] = nw.Appender()
+	return w.activeAppenders[hash]
+}
+
+func (w *MultiWriter) Add(l labels.Labels, t int64, v float64) (uint64, error) {
+	return w.getOrCreate(t).Add(l, t, v)
+}
+
+func (w *MultiWriter) AddFast(ref uint64, t int64, v float64) error {
+	return w.getOrCreate(t).AddFast(ref, t, v)
+}
+
+func (w *MultiWriter) Commit() error {
+	var merr tsdb_errors.MultiError
+	for _, a := range w.activeAppenders {
+		merr.Add(a.Commit())
+	}
+	return merr.Err()
+}
+
+func (w *MultiWriter) Rollback() error {
+	var merr tsdb_errors.MultiError
+	for _, a := range w.activeAppenders {
+		merr.Add(a.Rollback())
+	}
+	return merr.Err()
+}
+
+func (w *MultiWriter) Flush() ([]ulid.ULID, error) {
+	ids := make([]ulid.ULID, 0, len(w.blocks))
+	for _, b := range w.blocks {
+		id, err := b.Flush()
+		if err != nil {
+			return nil, err
+		}
+		ids = append(ids, id...)
+	}
+	return ids, nil
+}
+
+func (w *MultiWriter) Close() error {
+	var merr tsdb_errors.MultiError
+	for _, b := range w.blocks {
+		merr.Add(b.Close())
+	}
+	return merr.Err()
+}