github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/storage/debug.go (about)

     1  package storage
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"fmt"
     7  	"io"
     8  	"net/http"
     9  
    10  	"github.com/dgraph-io/badger/v2"
    11  	"github.com/gorilla/mux"
    12  )
    13  
    14  func (s *Storage) DebugExport(w http.ResponseWriter, r *http.Request) {
    15  	select {
    16  	default:
    17  	case <-s.stop:
    18  		w.WriteHeader(http.StatusServiceUnavailable)
    19  		return
    20  	}
    21  
    22  	n := mux.Vars(r)["db"]
    23  	var d BadgerDBWithCache
    24  	switch n {
    25  	case "segments":
    26  		d = s.segments
    27  	case "trees":
    28  		d = s.trees
    29  	case "dicts":
    30  		d = s.dicts
    31  	case "dimensions":
    32  		d = s.dimensions
    33  	default:
    34  		// Note that export from main DB is not allowed.
    35  		http.Error(w, fmt.Sprintf("database %q not found", n), http.StatusNotFound)
    36  		return
    37  	}
    38  
    39  	// TODO(kolesnikovae): Refactor routes registration and use gorilla Queries.
    40  	k, ok := r.URL.Query()["k"]
    41  	if !ok || len(k) != 1 {
    42  		http.Error(w, "query parameter 'k' is required", http.StatusBadRequest)
    43  		return
    44  	}
    45  
    46  	w.Header().Set("Content-Type", "application/octet-stream")
    47  	err := d.View(func(txn *badger.Txn) error {
    48  		item, err := txn.Get([]byte(k[0]))
    49  		if err != nil {
    50  			return err
    51  		}
    52  		return item.Value(func(v []byte) error {
    53  			_, err = io.Copy(w, bytes.NewBuffer(v))
    54  			return err
    55  		})
    56  	})
    57  
    58  	switch {
    59  	case err == nil:
    60  	case errors.Is(err, badger.ErrKeyNotFound):
    61  		http.Error(w, fmt.Sprintf("key %q not found in %s", k[0], n), http.StatusNotFound)
    62  	default:
    63  		http.Error(w, fmt.Sprintf("failed to export value for key %q: %v", k[0], err), http.StatusInternalServerError)
    64  	}
    65  }