Added query blocks cli command (#876)

* Added query blocks cli command Signed-off-by: Joe Elliott <[email protected]> * changelog Signed-off-by: Joe Elliott <[email protected]> * docs Signed-off-by: Joe Elliott <[email protected]> * intensity in ten cities! Signed-off-by: Joe Elliott <[email protected]>
grafana · Aug 13, 2021 · 5bd2fc6 · 5bd2fc6
1 parent f822caa
commit 5bd2fc6
Show file tree

Hide file tree

Showing 6 changed files with 192 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## main / unreleased
 
 * [BUGFIX] Update port spec for GCS docker-compose example [#869](https://github.com/grafana/tempo/pull/869) (@zalegrala)
+* [ENHANCEMENT] Added "query blocks" cli option. [#876](https://github.com/grafana/tempo/pull/876) (@joe-elliott)
 
 ## v1.1.0-rc.0 / 2021-08-11
 

diff --git a/cmd/tempo-cli/cmd-query-blocks.go b/cmd/tempo-cli/cmd-query-blocks.go
@@ -0,0 +1,162 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strconv"
+
+	"github.com/google/uuid"
+	"github.com/grafana/tempo/pkg/boundedwaitgroup"
+	"github.com/grafana/tempo/pkg/model"
+	"github.com/grafana/tempo/pkg/tempopb"
+	"github.com/grafana/tempo/pkg/util"
+	"github.com/grafana/tempo/tempodb/backend"
+	"github.com/grafana/tempo/tempodb/encoding"
+	"github.com/grafana/tempo/tempodb/encoding/common"
+)
+
+type queryResults struct {
+	blockID uuid.UUID
+	trace   *tempopb.Trace
+}
+
+type queryBlocksCmd struct {
+	backendOptions
+
+	TraceID  string `arg:"" help:"trace ID to retrieve"`
+	TenantID string `arg:"" help:"tenant ID to search"`
+}
+
+func (cmd *queryBlocksCmd) Run(ctx *globalOptions) error {
+	r, c, err := loadBackend(&cmd.backendOptions, ctx)
+	if err != nil {
+		return err
+	}
+
+	id, err := util.HexStringToTraceID(cmd.TraceID)
+	if err != nil {
+		return err
+	}
+
+	results, err := queryBucket(context.Background(), r, c, cmd.TenantID, id)
+	if err != nil {
+		return err
+	}
+
+	var combinedTrace *tempopb.Trace
+
+	fmt.Println()
+	for _, result := range results {
+		fmt.Println(result.blockID, ":")
+
+		jsonBytes, err := json.Marshal(result.trace)
+		if err != nil {
+			fmt.Println("failed to marshal to json: ", err)
+			continue
+		}
+
+		fmt.Println(string(jsonBytes))
+		combinedTrace, _, _, _ = model.CombineTraceProtos(result.trace, combinedTrace)
+	}
+
+	fmt.Println("combined:")
+	jsonBytes, err := json.Marshal(combinedTrace)
+	if err != nil {
+		fmt.Println("failed to marshal to json: ", err)
+		return nil
+	}
+	fmt.Println(string(jsonBytes))
+	return nil
+}
+
+func queryBucket(ctx context.Context, r backend.Reader, c backend.Compactor, tenantID string, traceID common.ID) ([]queryResults, error) {
+	blockIDs, err := r.Blocks(context.Background(), tenantID)
+	if err != nil {
+		return nil, err
+	}
+
+	fmt.Println("total blocks: ", len(blockIDs))
+
+	// Load in parallel
+	wg := boundedwaitgroup.New(20)
+	resultsCh := make(chan queryResults, len(blockIDs))
+
+	for blockNum, id := range blockIDs {
+		wg.Add(1)
+
+		go func(blockNum2 int, id2 uuid.UUID) {
+			defer wg.Done()
+
+			// search here
+			q, err := queryBlock(ctx, r, c, blockNum2, id2, tenantID, traceID)
+			if err != nil {
+				fmt.Println("Error querying block:", err)
+				return
+			}
+
+			if q != nil {
+				resultsCh <- *q
+			}
+		}(blockNum, id)
+	}
+
+	wg.Wait()
+	close(resultsCh)
+
+	results := make([]queryResults, 0)
+	for q := range resultsCh {
+		results = append(results, q)
+	}
+
+	return results, nil
+}
+
+func queryBlock(ctx context.Context, r backend.Reader, c backend.Compactor, blockNum int, id uuid.UUID, tenantID string, traceID common.ID) (*queryResults, error) {
+	fmt.Print(".")
+	if blockNum%100 == 0 {
+		fmt.Print(strconv.Itoa(blockNum))
+	}
+
+	meta, err := r.BlockMeta(context.Background(), id, tenantID)
+	if err != nil && err != backend.ErrDoesNotExist {
+		return nil, err
+	}
+
+	if err == backend.ErrDoesNotExist {
+		compactedMeta, err := c.CompactedBlockMeta(id, tenantID)
+		if err != nil && err != backend.ErrDoesNotExist {
+			return nil, err
+		}
+
+		if compactedMeta == nil {
+			return nil, fmt.Errorf("compacted meta nil?")
+		}
+
+		meta = &compactedMeta.BlockMeta
+	}
+
+	block, err := encoding.NewBackendBlock(meta, r)
+	if err != nil {
+		return nil, err
+	}
+
+	obj, err := block.Find(ctx, traceID)
+	if err != nil {
+		return nil, err
+	}
+
+	if obj == nil {
+		return nil, nil
+	}
+
+	trace, err := model.Unmarshal(obj, meta.DataEncoding)
+	if err != nil {
+		return nil, err
+	}
+
+	return &queryResults{
+		blockID: id,
+		trace:   trace,
+	}, nil
+}
diff --git a/cmd/tempo-cli/main.go b/cmd/tempo-cli/main.go
@@ -44,7 +44,10 @@ var cli struct {
 		Index viewIndexCmd `cmd:"" help:"View contents of block index"`
 	} `cmd:""`
 
-	Query queryCmd `cmd:"" help:"query tempo api"`
+	Query struct {
+		API    queryCmd       `cmd:"" help:"query tempo http api"`
+		Blocks queryBlocksCmd `cmd:"" help:"query for a traceid directly from backend blocks"`
+	} `cmd:""`
 }
 
 func main() {

diff --git a/docs/tempo/website/operations/tempo_cli.md b/docs/tempo/website/operations/tempo_cli.md
@@ -48,10 +48,10 @@ The backend can be configured in a few ways:
 
 Each option applies only to the command in which it is used. For example, `--backend <value>` does not permanently change where Tempo stores data. It only changes it for command in which you apply the option.
 
-## Query Command
+## Query API Command
 Call the tempo API and retrieve a trace by ID.
 ```bash
-tempo-cli query <api-endpoint> <trace-id>
+tempo-cli query api <api-endpoint> <trace-id>
 ```
 
 Arguments:
@@ -63,7 +63,26 @@ Options:
 
 **Example:**
 ```bash
-tempo-cli query http://tempo:3200 f1cfe82a8eef933b
+tempo-cli query api http://tempo:3200 f1cfe82a8eef933b
+```
+
+## Query Blocks Command
+Iterate over all backend blocks and dump all data found for a given trace id.
+```bash
+tempo-cli query blocks <trace-id> <tenant-id>
+```
+ **Note:** can be intense as it downloads every bloom filter and some percentage of indexes/trace data.
+
+Arguments:
+- `trace-id` Trace ID as a hexadecimal string.
+- `tenant-id` Tenant to search.
+
+Options:
+See backend options above.
+
+**Example:**
+```bash
+tempo-cli query blocks f1cfe82a8eef933b single-tenant
 ```
 
 ## List Blocks

diff --git a/pkg/util/http.go b/pkg/util/http.go
@@ -24,15 +24,15 @@ func ParseTraceID(r *http.Request) ([]byte, error) {
 		return nil, fmt.Errorf("please provide a traceID")
 	}
 
-	byteID, err := hexStringToTraceID(traceID)
+	byteID, err := HexStringToTraceID(traceID)
 	if err != nil {
 		return nil, err
 	}
 
 	return byteID, nil
 }
 
-func hexStringToTraceID(id string) ([]byte, error) {
+func HexStringToTraceID(id string) ([]byte, error) {
 	// The encoding/hex package does not handle non-hex characters.
 	// Ensure the ID has only the proper characters
 	for pos, idChar := range strings.Split(id, "") {

diff --git a/pkg/util/http_test.go b/pkg/util/http_test.go
@@ -46,7 +46,7 @@ func TestHexStringToTraceID(t *testing.T) {
 
 	for _, tt := range tc {
 		t.Run(tt.id, func(t *testing.T) {
-			actual, err := hexStringToTraceID(tt.id)
+			actual, err := HexStringToTraceID(tt.id)
 
 			if tt.expectError != nil {
 				assert.Equal(t, tt.expectError, err)