From a8772a41782758e8153c494dcd207e770d8421a4 Mon Sep 17 00:00:00 2001
From: Nidhey Nitin Indurkar <46122307+nidhey27@users.noreply.github.com>
Date: Thu, 1 Jun 2023 17:13:09 +0530
Subject: [PATCH] Feat: Get block by id directly on promtool analyze & get
 latest block if ID not provided (#12031)

* feat: analyze latest block or block by ID in CLI (promtool)

Signed-off-by: nidhey27 <nidhey.indurkar@infracloud.io>

* address remarks

Signed-off-by: nidhey60@gmail.com <nidhey.indurkar@infracloud.io>

* address latest review comments

Signed-off-by: nidhey60@gmail.com <nidhey.indurkar@infracloud.io>

---------

Signed-off-by: nidhey27 <nidhey.indurkar@infracloud.io>
Signed-off-by: nidhey60@gmail.com <nidhey.indurkar@infracloud.io>
---
 cmd/promtool/tsdb.go | 28 +++++++++--------------
 tsdb/db.go           | 54 ++++++++++++++++++++++++++++++++++++++++++++
 tsdb/db_test.go      | 19 +++++++++++++++-
 3 files changed, 83 insertions(+), 18 deletions(-)

diff --git a/cmd/promtool/tsdb.go b/cmd/promtool/tsdb.go
index 84aa43a9c..b7fad5fe0 100644
--- a/cmd/promtool/tsdb.go
+++ b/cmd/promtool/tsdb.go
@@ -398,26 +398,20 @@ func openBlock(path, blockID string) (*tsdb.DBReadOnly, tsdb.BlockReader, error)
 	if err != nil {
 		return nil, nil, err
 	}
-	blocks, err := db.Blocks()
-	if err != nil {
-		return nil, nil, err
-	}
-	var block tsdb.BlockReader
-	switch {
-	case blockID != "":
-		for _, b := range blocks {
-			if b.Meta().ULID.String() == blockID {
-				block = b
-				break
-			}
+
+	if blockID == "" {
+		blockID, err = db.LastBlockID()
+		if err != nil {
+			return nil, nil, err
 		}
-	case len(blocks) > 0:
-		block = blocks[len(blocks)-1]
 	}
-	if block == nil {
-		return nil, nil, fmt.Errorf("block %s not found", blockID)
+
+	b, err := db.Block(blockID)
+	if err != nil {
+		return nil, nil, err
 	}
-	return db, block, nil
+
+	return db, b, nil
 }
 
 func analyzeBlock(path, blockID string, limit int, runExtended bool) error {
diff --git a/tsdb/db.go b/tsdb/db.go
index 12974150b..32dae57a5 100644
--- a/tsdb/db.go
+++ b/tsdb/db.go
@@ -607,6 +607,60 @@ func (db *DBReadOnly) Blocks() ([]BlockReader, error) {
 	return blockReaders, nil
 }
 
+// LastBlockID returns the BlockID of latest block.
+func (db *DBReadOnly) LastBlockID() (string, error) {
+	entries, err := os.ReadDir(db.dir)
+	if err != nil {
+		return "", err
+	}
+
+	max := uint64(0)
+
+	lastBlockID := ""
+
+	for _, e := range entries {
+		// Check if dir is a block dir or not.
+		dirName := e.Name()
+		ulidObj, err := ulid.ParseStrict(dirName)
+		if err != nil {
+			continue // Not a block dir.
+		}
+		timestamp := ulidObj.Time()
+		if timestamp > max {
+			max = timestamp
+			lastBlockID = dirName
+		}
+	}
+
+	if lastBlockID == "" {
+		return "", errors.New("no blocks found")
+	}
+
+	return lastBlockID, nil
+}
+
+// Block returns a block reader by given block id.
+func (db *DBReadOnly) Block(blockID string) (BlockReader, error) {
+	select {
+	case <-db.closed:
+		return nil, ErrClosed
+	default:
+	}
+
+	_, err := os.Stat(filepath.Join(db.dir, blockID))
+	if os.IsNotExist(err) {
+		return nil, errors.Errorf("invalid block ID %s", blockID)
+	}
+
+	block, err := OpenBlock(db.logger, filepath.Join(db.dir, blockID), nil)
+	if err != nil {
+		return nil, err
+	}
+	db.closers = append(db.closers, block)
+
+	return block, nil
+}
+
 // Close all block readers.
 func (db *DBReadOnly) Close() error {
 	select {
diff --git a/tsdb/db_test.go b/tsdb/db_test.go
index 0b980f6ec..427a3b7af 100644
--- a/tsdb/db_test.go
+++ b/tsdb/db_test.go
@@ -2384,6 +2384,7 @@ func TestDBReadOnly(t *testing.T) {
 		dbDir     string
 		logger    = log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
 		expBlocks []*Block
+		expBlock  *Block
 		expSeries map[string][]tsdbutil.Sample
 		expChunks map[string][][]tsdbutil.Sample
 		expDBHash []byte
@@ -2427,6 +2428,7 @@ func TestDBReadOnly(t *testing.T) {
 		require.NoError(t, app.Commit())
 
 		expBlocks = dbWritable.Blocks()
+		expBlock = expBlocks[0]
 		expDbSize, err := fileutil.DirSize(dbWritable.Dir())
 		require.NoError(t, err)
 		require.Greater(t, expDbSize, dbSizeBeforeAppend, "db size didn't increase after an append")
@@ -2455,7 +2457,22 @@ func TestDBReadOnly(t *testing.T) {
 			require.Equal(t, expBlock.Meta(), blocks[i].Meta(), "block meta mismatch")
 		}
 	})
-
+	t.Run("block", func(t *testing.T) {
+		blockID := expBlock.meta.ULID.String()
+		block, err := dbReadOnly.Block(blockID)
+		require.NoError(t, err)
+		require.Equal(t, expBlock.Meta(), block.Meta(), "block meta mismatch")
+	})
+	t.Run("invalid block ID", func(t *testing.T) {
+		blockID := "01GTDVZZF52NSWB5SXQF0P2PGF"
+		_, err := dbReadOnly.Block(blockID)
+		require.Error(t, err)
+	})
+	t.Run("last block ID", func(t *testing.T) {
+		blockID, err := dbReadOnly.LastBlockID()
+		require.NoError(t, err)
+		require.Equal(t, expBlocks[2].Meta().ULID.String(), blockID)
+	})
 	t.Run("querier", func(t *testing.T) {
 		// Open a read only db and ensure that the API returns the same result as the normal DB.
 		q, err := dbReadOnly.Querier(context.TODO(), math.MinInt64, math.MaxInt64)