1// Copyright (c) 2017 Couchbase, Inc. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package cmd 16 17import ( 18 "encoding/csv" 19 "fmt" 20 "io" 21 "log" 22 "os" 23 "strconv" 24 25 "github.com/couchbase/vellum" 26 "github.com/spf13/cobra" 27) 28 29var mapCmd = &cobra.Command{ 30 Use: "map", 31 Short: "Map builds a new FST from a CSV file containing key,val pairs", 32 Long: `Map builds a new FST from a CSV file containing key,val pairs.`, 33 PreRunE: func(cmd *cobra.Command, args []string) error { 34 if len(args) < 1 { 35 return fmt.Errorf("source and target paths are required") 36 } 37 if len(args) < 2 { 38 return fmt.Errorf("target path is required") 39 } 40 return nil 41 }, 42 RunE: func(cmd *cobra.Command, args []string) error { 43 44 if !sorted { 45 return fmt.Errorf("only sorted input supported at this time") 46 } 47 48 file, err := os.Open(args[0]) 49 if err != nil { 50 log.Fatal(err) 51 } 52 defer file.Close() 53 54 f, err := os.Create(args[1]) 55 if err != nil { 56 return err 57 } 58 59 b, err := vellum.New(f, nil) 60 if err != nil { 61 return err 62 } 63 64 reader := csv.NewReader(file) 65 reader.FieldsPerRecord = 2 66 67 var record []string 68 record, err = reader.Read() 69 for err == nil { 70 var v uint64 71 v, err = strconv.ParseUint(record[1], 10, 64) 72 if err != nil { 73 return err 74 } 75 err = b.Insert([]byte(record[0]), v) 76 if err != nil { 77 return err 78 } 79 80 record, err = reader.Read() 81 } 82 if err != io.EOF { 83 return err 84 } 85 86 err = b.Close() 87 if err != nil { 88 return err 89 } 90 91 return nil 92 }, 93} 94 95func init() { 96 RootCmd.AddCommand(mapCmd) 97 mapCmd.Flags().BoolVar(&sorted, "sorted", false, "input already sorted") 98} 99