fork of https://github.com/sourcegraph/zoekt
1// Licensed under the Apache License, Version 2.0 (the "License");
2// you may not use this file except in compliance with the License.
3// You may obtain a copy of the License at
4//
5// http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS,
9// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10// See the License for the specific language governing permissions and
11// limitations under the License.
12
13// Command zoekt-merge-index merges a set of index shards into a compound shard.
14package main
15
16import (
17 "bufio"
18 "flag"
19 "fmt"
20 "log"
21 "os"
22 "path/filepath"
23 "strings"
24
25 "github.com/sourcegraph/zoekt/index"
26)
27
28// merge merges the input shards into a compound shard in dstDir. It returns the
29// full path to the compound shard. The input shards are removed on success.
30func merge(dstDir string, names []string) (string, error) {
31 var files []index.IndexFile
32 for _, fn := range names {
33 f, err := os.Open(fn)
34 if err != nil {
35 return "", nil
36 }
37 defer f.Close()
38
39 indexFile, err := index.NewIndexFile(f)
40 if err != nil {
41 return "", err
42 }
43 defer indexFile.Close()
44
45 files = append(files, indexFile)
46 }
47
48 tmpName, dstName, err := index.Merge(dstDir, files...)
49 if err != nil {
50 return "", err
51 }
52
53 // Delete input shards.
54 for _, name := range names {
55 paths, err := index.IndexFilePaths(name)
56 if err != nil {
57 return "", fmt.Errorf("zoekt-merge-index: %w", err)
58 }
59 for _, p := range paths {
60 if err := os.Remove(p); err != nil {
61 return "", fmt.Errorf("zoekt-merge-index: failed to remove simple shard: %w", err)
62 }
63 }
64 }
65
66 // We only rename the compound shard if all simple shards could be deleted in the
67 // previous step. This guarantees we won't have duplicate indexes.
68 if err := os.Rename(tmpName, dstName); err != nil {
69 return "", fmt.Errorf("zoekt-merge-index: failed to rename compound shard: %w", err)
70 }
71
72 return dstName, nil
73}
74
75func mergeCmd(paths []string) (string, error) {
76 if len(paths) == 0 {
77 return "", fmt.Errorf("merge requires at least one shard path")
78 }
79
80 if paths[0] == "-" {
81 paths = []string{}
82 scanner := bufio.NewScanner(os.Stdin)
83 for scanner.Scan() {
84 paths = append(paths, strings.TrimSpace(scanner.Text()))
85 }
86 if err := scanner.Err(); err != nil {
87 return "", err
88 }
89 log.Printf("merging %d paths from stdin", len(paths))
90 if len(paths) == 0 {
91 return "", fmt.Errorf("merge requires at least one shard path")
92 }
93 }
94
95 return merge(filepath.Dir(paths[0]), paths)
96}
97
98func explodeCmd(path string) error {
99 return index.Explode(filepath.Dir(path), path)
100}
101
102func main() {
103 flag.Usage = func() {
104 fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s <command> [args]\n\n", os.Args[0])
105 fmt.Fprintln(flag.CommandLine.Output(), "Commands:")
106 fmt.Fprintln(flag.CommandLine.Output(), " merge <shard> [<shard>...]")
107 fmt.Fprintln(flag.CommandLine.Output(), " merge -")
108 fmt.Fprintln(flag.CommandLine.Output(), " explode <compound-shard>")
109 }
110 flag.Parse()
111
112 switch subCommand := flag.Arg(0); subCommand {
113 case "merge":
114 compoundShardPath, err := mergeCmd(flag.Args()[1:])
115 if err != nil {
116 log.Fatal(err)
117 }
118 fmt.Println(compoundShardPath)
119 case "explode":
120 if flag.NArg() != 2 {
121 log.Fatal("explode requires exactly one compound shard path")
122 }
123 if err := explodeCmd(flag.Arg(1)); err != nil {
124 log.Fatal(err)
125 }
126 default:
127 fmt.Fprintf(os.Stderr, "unknown subcommand %q\n", subCommand)
128 flag.Usage()
129 os.Exit(2)
130 }
131}