From e2dd5b61efa58ac8adc0eb1fb6291af7ef059c62 Mon Sep 17 00:00:00 2001
From: Dipack P Panjabi <dipack.panjabi@gmail.com>
Date: Thu, 21 Nov 2019 07:10:25 -0500
Subject: [PATCH] Added CreateBlock and CreateHead functions to new file
 (#6331)

* Added CreateBlock and CreateHead functions to new file to make it reusable across packages.

Signed-off-by: Dipack P Panjabi <dipack.panjabi@gmail.com>
---
 tsdb/db.go            |  7 +++-
 tsdb/tsdbblockutil.go | 83 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 tsdb/tsdbblockutil.go

diff --git a/tsdb/db.go b/tsdb/db.go
index 9ab34767b..715a97b28 100644
--- a/tsdb/db.go
+++ b/tsdb/db.go
@@ -43,12 +43,17 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+// Default duration of a block in milliseconds - 2h.
+const (
+	DefaultBlockDuration = int64(2 * 60 * 60 * 1000)
+)
+
 // DefaultOptions used for the DB. They are sane for setups using
 // millisecond precision timestamps.
 var DefaultOptions = &Options{
 	WALSegmentSize:         wal.DefaultSegmentSize,
 	RetentionDuration:      15 * 24 * 60 * 60 * 1000, // 15 days in milliseconds
-	BlockRanges:            ExponentialBlockRanges(int64(2*time.Hour)/1e6, 3, 5),
+	BlockRanges:            ExponentialBlockRanges(DefaultBlockDuration, 3, 5),
 	NoLockfile:             false,
 	AllowOverlappingBlocks: false,
 	WALCompression:         false,
diff --git a/tsdb/tsdbblockutil.go b/tsdb/tsdbblockutil.go
new file mode 100644
index 000000000..eb1756019
--- /dev/null
+++ b/tsdb/tsdbblockutil.go
@@ -0,0 +1,83 @@
+// Copyright 2019 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tsdb
+
+import (
+	"context"
+	"fmt"
+	"github.com/go-kit/kit/log"
+	"github.com/prometheus/prometheus/pkg/labels"
+	"os"
+	"path/filepath"
+)
+
+var InvalidTimesError = fmt.Errorf("max time is lesser than min time")
+
+type MetricSample struct {
+	TimestampMs int64
+	Value       float64
+	Labels      labels.Labels
+}
+
+// CreateHead creates a TSDB writer head to write the sample data to.
+func CreateHead(samples []*MetricSample, chunkRange int64, logger log.Logger) (*Head, error) {
+	head, err := NewHead(nil, logger, nil, chunkRange)
+	if err != nil {
+		return nil, err
+	}
+	app := head.Appender()
+	for _, sample := range samples {
+		_, err = app.Add(sample.Labels, sample.TimestampMs, sample.Value)
+		if err != nil {
+			return nil, err
+		}
+	}
+	err = app.Commit()
+	if err != nil {
+		return nil, err
+	}
+	return head, nil
+}
+
+// CreateBlock creates a chunkrange block from the samples passed to it, and writes it to disk.
+func CreateBlock(samples []*MetricSample, dir string, mint, maxt int64, logger log.Logger) (string, error) {
+	chunkRange := maxt - mint
+	if chunkRange == 0 {
+		chunkRange = DefaultBlockDuration
+	}
+	if chunkRange < 0 {
+		return "", InvalidTimesError
+	}
+	head, err := CreateHead(samples, chunkRange, logger)
+	if err != nil {
+		return "", err
+	}
+
+	compactor, err := NewLeveledCompactor(context.Background(), nil, logger, DefaultOptions.BlockRanges, nil)
+	if err != nil {
+		return "", err
+	}
+
+	err = os.MkdirAll(dir, 0777)
+	if err != nil {
+		return "", err
+	}
+
+	ulid, err := compactor.Write(dir, head, mint, maxt, nil)
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(dir, ulid.String()), nil
+}