add tsmdir cmd line option to inspect

influxdata · Aug 4, 2016 · 2ccf4bd · 2ccf4bd
1 parent 6e60cbb
commit 2ccf4bd
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 14 deletions.
diff --git a/cmd/influx_inspect/export.go b/cmd/influx_inspect/export.go
@@ -4,6 +4,7 @@ import (
 	"compress/gzip"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
 	"os"
 	"path/filepath"
@@ -15,6 +16,7 @@ import (
 
 type cmdExport struct {
 	path            string
+	tsmDir          string
 	out             string
 	db              string
 	retentionPolicy string
@@ -24,9 +26,10 @@ type cmdExport struct {
 	files map[string][]string
 }
 
-func newCmdExport(path, out, db, retentionPolicy string, compress bool) *cmdExport {
+func newCmdExport(path, tsmDir, out, db, retentionPolicy string, compress bool) *cmdExport {
 	return &cmdExport{
 		path:            filepath.Join(path, "data"),
+		tsmDir:          tsmDir,
 		out:             out,
 		db:              db,
 		compress:        compress,
@@ -53,7 +56,11 @@ func (c *cmdExport) run() error {
 }
 
 func (c *cmdExport) export() error {
-	if err := c.walkFiles(); err != nil {
+	if c.tsmDir != "" {
+		if err := c.getTSMDirFiles(); err != nil {
+			return err
+		}
+	} else if err := c.walkFiles(); err != nil {
 		return err
 	}
 	return c.writeFiles()
@@ -99,19 +106,25 @@ func (c *cmdExport) writeFiles() error {
 		w = gzip.NewWriter(w)
 	}
 
-	// Write out all the DDL
-	fmt.Fprintln(w, "# DDL")
-	for key, _ := range c.files {
-		keys := strings.Split(key, string(byte(os.PathSeparator)))
-		fmt.Fprintf(w, "CREATE DATABASE %s\n", keys[0])
-		fmt.Fprintf(w, "CREATE RETENTION POLICY %s ON %s DURATION inf REPLICATION 1\n", keys[1], keys[0])
+	if c.tsmDir == "" {
+		// Write out all the DDL
+		fmt.Fprintln(w, "# DDL")
+		for key, _ := range c.files {
+			keys := strings.Split(key, string(byte(os.PathSeparator)))
+			fmt.Fprintf(w, "CREATE DATABASE %s\n", keys[0])
+			fmt.Fprintf(w, "CREATE RETENTION POLICY %s ON %s DURATION inf REPLICATION 1\n", keys[1], keys[0])
+		}
+
+		fmt.Fprintln(w, "# DML")
 	}
 
-	fmt.Fprintln(w, "# DML")
 	for key, files := range c.files {
-		keys := strings.Split(key, string(byte(os.PathSeparator)))
-		fmt.Fprintf(w, "# CONTEXT-DATABASE:%s\n", keys[0])
-		fmt.Fprintf(w, "# CONTEXT-RETENTION-POLICY:%s\n", keys[1])
+		if c.tsmDir == "" {
+			keys := strings.Split(key, string(byte(os.PathSeparator)))
+			fmt.Fprintf(w, "# CONTEXT-DATABASE:%s\n", keys[0])
+			fmt.Fprintf(w, "# CONTEXT-RETENTION-POLICY:%s\n", keys[1])
+		}
+
 		for _, f := range files {
 			// use an anonymous function here to close the files in the defers and not let them
 			// accumulate in the loop
@@ -160,3 +173,28 @@ func (c *cmdExport) writeFiles() error {
 	}
 	return nil
 }
+
+func (c *cmdExport) getTSMDirFiles() error {
+	fis, err := ioutil.ReadDir(c.tsmDir)
+	if err != nil {
+		return err
+	}
+
+	files := []string{}
+
+	for _, fi := range fis {
+		if fi.IsDir() {
+			continue
+		}
+
+		ext := filepath.Ext(fi.Name())
+
+		if ext == c.ext {
+			files = append(files, filepath.Join(c.tsmDir, fi.Name()))
+		}
+	}
+
+	c.files[""] = files
+
+	return nil
+}
diff --git a/cmd/influx_inspect/main.go b/cmd/influx_inspect/main.go
@@ -129,10 +129,11 @@ func main() {
 		opts.path = fs.Args()[0]
 		cmdExportWAL(opts)
 	case "export":
-		var path, out, db, rp string
+		var path, tsmDir, out, db, rp string
 		var compress bool
 		fs := flag.NewFlagSet("export", flag.ExitOnError)
 		fs.StringVar(&path, "dir", os.Getenv("HOME")+"/.influxdb", "Root storage path. [$HOME/.influxdb]")
+		fs.StringVar(&tsmDir, "tsmdir", "", "Path to single dir with tsm files to export.")
 		fs.StringVar(&out, "out", os.Getenv("HOME")+"/.influxdb/export", "Destination file to export to")
 		fs.StringVar(&db, "db", "", "Optional: the database to export")
 		fs.StringVar(&rp, "rp", "", "Optional: the retention policy to export (requires db parameter to be specified)")
@@ -149,7 +150,7 @@ func main() {
 			fmt.Printf("%v", err)
 			os.Exit(1)
 		}
-		c := newCmdExport(path, out, db, rp, compress)
+		c := newCmdExport(path, tsmDir, out, db, rp, compress)
 		if err := c.run(); err != nil {
 			fmt.Println(err)
 			os.Exit(1)