From 691177e68ade47034515339b90bbfffc53dc68b6 Mon Sep 17 00:00:00 2001
From: Alex R <i@sepa.spb.ru>
Date: Sat, 18 May 2024 16:17:30 +0200
Subject: [PATCH] unwrap: add relabel for meta.json

---
 main.go   |  3 ++-
 unwrap.go | 19 ++++++++++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index cb89b96..7c02fce 100644
--- a/main.go
+++ b/main.go
@@ -58,6 +58,7 @@ func main() {
 
 	unwrapCmd := app.Command("unwrap", "Split TSDB block to multiple blocks by Label")
 	unwrapRelabel := extkingpin.RegisterPathOrContent(unwrapCmd, "relabel-config", fmt.Sprintf("YAML file that contains relabeling configuration. Set %s=name1;name2;... to split separate blocks for each uniq label combination.", metaExtLabels), extkingpin.WithEnvSubstitution(), extkingpin.WithRequired())
+	unwrapMetaRelabel := extkingpin.RegisterPathOrContent(unwrapCmd, "meta-relabel", "YAML file that contains relabeling configuration for block labels (meta.json)", extkingpin.WithEnvSubstitution())
 	unwrapRecursive := unwrapCmd.Flag("recursive", "Recursive search for blocks in the bucket (Mimir has blocks nested to tenants folders)").Short('r').Default("false").Bool()
 	unwrapDir := unwrapCmd.Flag("data-dir", "Data directory in which to cache blocks and process tsdb.").Default("./data").String()
 	unwrapWait := unwrapCmd.Flag("wait-interval", "Wait interval between consecutive runs and bucket refreshes. Run once if 0.").Default("5m").Short('w').Duration()
@@ -100,7 +101,7 @@ func main() {
 	case importCmd.FullCommand():
 		exitCode(importMetrics(bkt, importFromFile, importBlockSize, importDir, importLabels, *importUpload, logger))
 	case unwrapCmd.FullCommand():
-		exitCode(unwrap(bkt, *unwrapRelabel, *unwrapRecursive, unwrapDir, unwrapWait, *unwrapDry, unwrapDst, unwrapMaxTime, unwrapSrc, logger))
+		exitCode(unwrap(bkt, *unwrapRelabel, *unwrapMetaRelabel, *unwrapRecursive, unwrapDir, unwrapWait, *unwrapDry, unwrapDst, unwrapMaxTime, unwrapSrc, logger))
 	}
 }
 
diff --git a/unwrap.go b/unwrap.go
index aa17889..7cbd7a3 100644
--- a/unwrap.go
+++ b/unwrap.go
@@ -32,7 +32,7 @@ import (
 
 const metaExtLabels = "__meta_ext_labels"
 
-func unwrap(bkt objstore.Bucket, unwrapRelabel extkingpin.PathOrContent, recursive bool, dir *string, wait *time.Duration, unwrapDry bool, outConfig *extkingpin.PathOrContent, maxTime *model.TimeOrDurationValue, unwrapSrc *string, logger log.Logger) (err error) {
+func unwrap(bkt objstore.Bucket, unwrapRelabel extkingpin.PathOrContent, unwrapMetaRelabel extkingpin.PathOrContent, recursive bool, dir *string, wait *time.Duration, unwrapDry bool, outConfig *extkingpin.PathOrContent, maxTime *model.TimeOrDurationValue, unwrapSrc *string, logger log.Logger) (err error) {
 	relabelContentYaml, err := unwrapRelabel.Content()
 	if err != nil {
 		return fmt.Errorf("get content of relabel configuration: %w", err)
@@ -41,6 +41,14 @@ func unwrap(bkt objstore.Bucket, unwrapRelabel extkingpin.PathOrContent, recursi
 	if err := yaml.Unmarshal(relabelContentYaml, &relabelConfig); err != nil {
 		return fmt.Errorf("parse relabel configuration: %w", err)
 	}
+	metaRelabelContentYaml, err := unwrapMetaRelabel.Content()
+	if err != nil {
+		return fmt.Errorf("get content of meta-relabel configuration: %w", err)
+	}
+	var metaRelabel []*relabel.Config
+	if err := yaml.Unmarshal(metaRelabelContentYaml, &metaRelabel); err != nil {
+		return fmt.Errorf("parse relabel configuration: %w", err)
+	}
 
 	objStoreYaml, err := outConfig.Content()
 	if err != nil {
@@ -72,7 +80,7 @@ func unwrap(bkt objstore.Bucket, unwrapRelabel extkingpin.PathOrContent, recursi
 					continue
 				}
 			}
-			if err := unwrapBlock(bkt, b, relabelConfig, *dir, unwrapDry, dst, logger); err != nil {
+			if err := unwrapBlock(bkt, b, relabelConfig, metaRelabel, *dir, unwrapDry, dst, logger); err != nil {
 				return err
 			}
 		}
@@ -90,7 +98,7 @@ func unwrap(bkt objstore.Bucket, unwrapRelabel extkingpin.PathOrContent, recursi
 	})
 }
 
-func unwrapBlock(bkt objstore.Bucket, b Block, relabelConfig []*relabel.Config, dir string, unwrapDry bool, dst objstore.Bucket, logger log.Logger) error {
+func unwrapBlock(bkt objstore.Bucket, b Block, relabelConfig []*relabel.Config, metaRelabel []*relabel.Config, dir string, unwrapDry bool, dst objstore.Bucket, logger log.Logger) error {
 	if err := runutil.DeleteAll(dir); err != nil {
 		return fmt.Errorf("unable to cleanup cache folder %s: %w", dir, err)
 	}
@@ -114,6 +122,11 @@ func unwrapBlock(bkt objstore.Bucket, b Block, relabelConfig []*relabel.Config,
 	if err != nil {
 		return fmt.Errorf("fail to read meta.json for %s: %w", b.Id.String(), err)
 	}
+	lbls, keep := relabel.Process(labels.FromMap(origMeta.Thanos.Labels), metaRelabel...)
+	if !keep {
+		return nil
+	}
+	origMeta.Thanos.Labels = lbls.Map()
 	db, err := tsdb.OpenDBReadOnly(inDir, logger)
 	if err != nil {
 		return err