From 615f1cf6f1194905edecf2f9212df8359032feb8 Mon Sep 17 00:00:00 2001
From: nidhey27 <nidhey.indurkar@infracloud.io>
Date: Mon, 6 Mar 2023 16:18:27 +0530
Subject: [PATCH] feat: analyze latest block or block by ID in CLI (promtool)

Signed-off-by: nidhey27 <nidhey.indurkar@infracloud.io>
---
 cmd/promtool/tsdb.go | 25 ++++++++------
 tsdb/db.go           | 79 ++++++++++++++++++++++++++++++++++++++++++++
 tsdb/db_test.go      | 22 +++++++++++-
 3 files changed, 114 insertions(+), 12 deletions(-)

diff --git a/cmd/promtool/tsdb.go b/cmd/promtool/tsdb.go
index 0e0cdb863a6..a3f5d3500f2 100644
--- a/cmd/promtool/tsdb.go
+++ b/cmd/promtool/tsdb.go
@@ -398,21 +398,24 @@ 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
+
+	if blockID == "" {
+		blockID, err = db.LastBlockID(nil)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
 	if blockID != "" {
-		for _, b := range blocks {
-			if b.Meta().ULID.String() == blockID {
-				block = b
-				break
-			}
+		b, err := db.Block(nil, filepath.Join(blockID))
+		if err != nil {
+			return nil, nil, err
 		}
-	} else if len(blocks) > 0 {
-		block = blocks[len(blocks)-1]
+		block = b
 	}
+
 	if block == nil {
 		return nil, nil, fmt.Errorf("block %s not found", blockID)
 	}
diff --git a/tsdb/db.go b/tsdb/db.go
index 659251c3ca3..66dd21cb448 100644
--- a/tsdb/db.go
+++ b/tsdb/db.go
@@ -601,6 +601,85 @@ func (db *DBReadOnly) Blocks() ([]BlockReader, error) {
 	return blockReaders, nil
 }
 
+// LastBlockID returns the BlockID of latest block.
+func (db *DBReadOnly) LastBlockID(logger log.Logger) (string, error) {
+	select {
+	case <-db.closed:
+		return "", ErrClosed
+	default:
+	}
+	latestDirName, err := db.lastBlockDirName(logger)
+	if err != nil {
+		return "", err
+	}
+	// Open the latest block and get its ID
+	block, err := db.Block(logger, latestDirName)
+	if err != nil {
+		return "", err
+	}
+	return block.Meta().ULID.String(), nil
+}
+
+func (db *DBReadOnly) lastBlockDirName(logger log.Logger) (string, error) {
+	entries, err := os.ReadDir(db.dir)
+	if err != nil {
+		return "", err
+	}
+	skipNames := map[string]struct{}{
+		"index":          {},
+		"chunks_head":    {},
+		"lock":           {},
+		"queries.active": {},
+		"wal":            {},
+	}
+	max := uint64(0)
+
+	latestDirName := ""
+	// Walk the blocks directory and find the latest subdirectory
+	for _, e := range entries {
+		dirName := e.Name()
+
+		if _, skip := skipNames[dirName]; skip {
+			// Skip the  directory
+			continue
+		}
+
+		ulidObj, err := ulid.Parse(dirName)
+		if err != nil {
+			return "", err
+		}
+		timestamp := ulidObj.Time()
+
+		if timestamp > max {
+			max = timestamp
+			latestDirName = dirName
+		}
+	}
+
+	if latestDirName == "" {
+		return "", errors.New("no blocks found")
+	}
+
+	return latestDirName, nil
+}
+
+// Block returns a block reader by given block id.
+func (db *DBReadOnly) Block(logger log.Logger, blockID string) (BlockReader, error) {
+	select {
+	case <-db.closed:
+		return nil, ErrClosed
+	default:
+	}
+
+	block, err := OpenBlock(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 70639085e41..4be6f5fd5fd 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,25 @@ func TestDBReadOnly(t *testing.T) {
 			require.Equal(t, expBlock.Meta(), blocks[i].Meta(), "block meta mismatch")
 		}
 	})
-
+	t.Run("analyse-block-id", func(t *testing.T) {
+		blockID := expBlock.meta.ULID.String()
+		block, err := dbReadOnly.Block(nil, 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(nil, blockID)
+		require.Error(t, err)
+	})
+	t.Run("analyse-latest-block", func(t *testing.T) {
+		blockID, err := dbReadOnly.LastBlockID(nil)
+		expBlock = expBlocks[2]
+		require.NoError(t, err)
+		block, err := dbReadOnly.Block(nil, blockID)
+		require.NoError(t, err)
+		require.Equal(t, expBlock.Meta(), block.Meta(), "block meta mismatch")
+	})
 	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)