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 "fmt"
19 "log"
20 "os"
21 "path/filepath"
22 "strings"
23
24 "github.com/sourcegraph/zoekt/index"
25)
26
27// merge merges the input shards into a compound shard in dstDir. It returns the
28// full path to the compound shard. The input shards are removed on success.
29func merge(dstDir string, names []string) (string, error) {
30 var files []index.IndexFile
31 for _, fn := range names {
32 f, err := os.Open(fn)
33 if err != nil {
34 return "", nil
35 }
36 defer f.Close()
37
38 indexFile, err := index.NewIndexFile(f)
39 if err != nil {
40 return "", err
41 }
42 defer indexFile.Close()
43
44 files = append(files, indexFile)
45 }
46
47 tmpName, dstName, err := index.Merge(dstDir, files...)
48 if err != nil {
49 return "", err
50 }
51
52 // Delete input shards.
53 for _, name := range names {
54 paths, err := index.IndexFilePaths(name)
55 if err != nil {
56 return "", fmt.Errorf("zoekt-merge-index: %w", err)
57 }
58 for _, p := range paths {
59 if err := os.Remove(p); err != nil {
60 return "", fmt.Errorf("zoekt-merge-index: failed to remove simple shard: %w", err)
61 }
62 }
63 }
64
65 // We only rename the compound shard if all simple shards could be deleted in the
66 // previous step. This guarantees we won't have duplicate indexes.
67 if err := os.Rename(tmpName, dstName); err != nil {
68 return "", fmt.Errorf("zoekt-merge-index: failed to rename compound shard: %w", err)
69 }
70
71 return dstName, nil
72}
73
74func mergeCmd(paths []string) (string, error) {
75 if paths[0] == "-" {
76 paths = []string{}
77 scanner := bufio.NewScanner(os.Stdin)
78 for scanner.Scan() {
79 paths = append(paths, strings.TrimSpace(scanner.Text()))
80 }
81 if err := scanner.Err(); err != nil {
82 return "", err
83 }
84 log.Printf("merging %d paths from stdin", len(paths))
85 }
86
87 return merge(filepath.Dir(paths[0]), paths)
88}
89
90func explodeCmd(path string) error {
91 return index.Explode(filepath.Dir(path), path)
92}
93
94func main() {
95 switch subCommand := os.Args[1]; subCommand {
96 case "merge":
97 compoundShardPath, err := mergeCmd(os.Args[2:])
98 if err != nil {
99 log.Fatal(err)
100 }
101 fmt.Println(compoundShardPath)
102 case "explode":
103 if err := explodeCmd(os.Args[2]); err != nil {
104 log.Fatal(err)
105 }
106 default:
107 log.Fatalf("unknown subcommand %s", subCommand)
108 }
109}