Merge pull request grafana/phlare#301 from grafana/20221010_merge-into-profilecli

Merge firecli and parquet-tool into profilecli

Author: simonswine

Makefile

@@ -67,6 +67,7 @@ go/deps:
 go/bin:
 	mkdir -p ./bin
 	CGO_ENABLED=0 $(GO) build $(GO_FLAGS) -o bin/ ./cmd/fire
+	CGO_ENABLED=0 $(GO) build $(GO_FLAGS) -o bin/ ./cmd/profilecli
 
 .PHONY: go/lint
 go/lint: $(BIN)/golangci-lint

cmd/fire/Dockerfile

@@ -15,6 +15,7 @@ FROM alpine:3.16.2
 RUN apk add --no-cache ca-certificates
 
 COPY --from=build /src/fire/bin/fire /usr/bin/fire
+COPY --from=build /src/fire/bin/profilecli /usr/bin/profilecli
 COPY cmd/fire/fire.yaml /etc/fire/config.yaml
 
 RUN addgroup -g 10001 -S fire && \

cmd/firetool/blocks.go renamed to cmd/profilecli/blocks.go

@@ -37,7 +37,7 @@ func blocksList(ctx context.Context) error {
 		return err
 	}
 
-	table := tablewriter.NewWriter(os.Stdout)
+	table := tablewriter.NewWriter(output(ctx))
 	table.SetHeader([]string{"Block ID", "MinTime", "MaxTime", "Duration", "Index", "Profiles", "Stacktraces", "Locations", "Functions", "Strings"})
 	for _, blockInfo := range metas {
 		table.Append([]string{

cmd/firetool/main.go renamed to cmd/profilecli/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 
@@ -29,6 +30,7 @@ var (
 
 func main() {
 	ctx := firecontext.WithLogger(context.Background(), logger)
+	ctx = withOutput(ctx, os.Stdout)
 
 	app := kingpin.New(filepath.Base(os.Args[0]), "Tooling for Grafana Fire, the continuous profiling aggregation system.").UsageWriter(os.Stdout)
 	app.Version(version.Print("firetool"))
@@ -41,6 +43,10 @@ func main() {
 	blocksListCmd := blocksCmd.Command("list", "List blocks.")
 	blocksListCmd.Flag("restore-missing-meta", "").Default("false").BoolVar(&cfg.blocks.restoreMissingMeta)
 
+	parquetCmd := app.Command("parquet", "Operate on a Parquet file.")
+	parquetInspectCmd := parquetCmd.Command("inspect", "Inspect a parquet file's structure.")
+	parquetInspectFiles := parquetInspectCmd.Arg("file", "parquet file path").Required().ExistingFiles()
+
 	parsedCmd := kingpin.MustParse(app.Parse(os.Args[1:]))
 
 	if !cfg.verbose {
@@ -50,6 +56,12 @@ func main() {
 	switch parsedCmd {
 	case blocksListCmd.FullCommand():
 		os.Exit(checkError(blocksList(ctx)))
+	case parquetInspectCmd.FullCommand():
+		for _, file := range *parquetInspectFiles {
+			if err := parquetInspect(ctx, file); err != nil {
+				os.Exit(checkError(err))
+			}
+		}
 	}
 }
 
@@ -60,3 +72,20 @@ func checkError(err error) int {
 	}
 	return 0
 }
+
+type contextKey uint8
+
+const (
+	contextKeyOutput contextKey = iota
+)
+
+func withOutput(ctx context.Context, w io.Writer) context.Context {
+	return context.WithValue(ctx, contextKeyOutput, w)
+}
+
+func output(ctx context.Context) io.Writer {
+	if w, ok := ctx.Value(contextKeyOutput).(io.Writer); ok {
+		return w
+	}
+	return os.Stdout
+}

cmd/parquet-tool/main.go renamed to cmd/profilecli/parquet.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -10,33 +11,30 @@ import (
 	"github.com/segmentio/parquet-go"
 )
 
-func main() {
-	if len(os.Args) < 2 {
-		fmt.Println("Usage: parquet-tool <file>")
-		os.Exit(1)
-	}
-	f, err := os.Open(os.Args[1])
+func parquetInspect(ctx context.Context, path string) error {
+	f, err := os.Open(path)
 	if err != nil {
-		panic(err)
+		return err
 	}
 	defer f.Close()
 	stats, err := f.Stat()
 	if err != nil {
-		panic(err)
+		return err
 	}
 	pf, err := parquet.OpenFile(f, stats.Size())
 	if err != nil {
-		panic(err)
+		return err
 	}
-	fmt.Println("schema:", pf.Schema())
+	out := output(ctx)
+	fmt.Fprintln(out, "schema:", pf.Schema())
 	meta := pf.Metadata()
 	fmt.Println("Num Rows:", meta.NumRows)
 	for i, rg := range meta.RowGroups {
-		fmt.Println("\t Row group:", i)
-		fmt.Println("\t\t Row Count:", rg.NumRows)
-		fmt.Println("\t\t Row size:", humanize.Bytes(uint64(rg.TotalByteSize)))
-		fmt.Println("\t\t Columns:")
-		table := tablewriter.NewWriter(os.Stdout)
+		fmt.Fprintln(out, "\t Row group:", i)
+		fmt.Fprintln(out, "\t\t Row Count:", rg.NumRows)
+		fmt.Fprintln(out, "\t\t Row size:", humanize.Bytes(uint64(rg.TotalByteSize)))
+		fmt.Fprintln(out, "\t\t Columns:")
+		table := tablewriter.NewWriter(out)
 		table.SetHeader([]string{
 			"Col", "Type", "NumVal", "TotalCompressedSize", "TotalUncompressedSize", "Compression", "%", "PageCount", "AvgPageSize",
 		})
@@ -63,4 +61,6 @@ func main() {
 		}
 		table.Render()
 	}
+
+	return nil
 }