-
Notifications
You must be signed in to change notification settings - Fork 16
/
fetch.go
138 lines (114 loc) · 3.2 KB
/
fetch.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
// Copyright (C) 2015 Foursquare Labs Inc.
package main
import (
"encoding/json"
"fmt"
"log"
"net/http"
"strings"
"github.com/foursquare/quiver/hfile"
)
type CollectionSpecList struct {
Collections []SingleCollectionSpec
}
type SingleCollectionSpec struct {
Capacity int
Collection string
Function string
LockNamespace string
Partition int
Path string
Url string
Ondemand bool
}
func getCollectionConfig(args []string) []*hfile.CollectionConfig {
if Settings.configJsonUrl != "" {
if len(args) > 0 {
log.Fatalln("Only one of command-line collection specs or json config may be used.")
}
return ConfigsFromJsonUrl(Settings.configJsonUrl)
}
return ConfigsFromCommandline(args)
}
func ConfigsFromCommandline(args []string) []*hfile.CollectionConfig {
configs := make([]*hfile.CollectionConfig, len(args))
for i, pair := range args {
nameAndPath := strings.SplitN(pair, "=", 2)
if len(nameAndPath) != 2 {
log.Fatal("collections must be specified in the form 'name=path'")
}
name, sfunc, total, part := nameAndPath[0], "_", "1", "0"
parent := name
if details := strings.SplitN(name, "/", 4); len(details) == 4 {
parent = details[0]
sfunc, total, part = details[1], details[2], details[3]
name = fmt.Sprintf("%s/%s", parent, part)
}
loadMethod := hfile.CopiedToMem
if Settings.onDisk {
loadMethod = hfile.OnDisk
} else if Settings.mlock {
loadMethod = hfile.MemlockFile
}
configs[i] = &hfile.CollectionConfig{
Name: name,
SourcePath: nameAndPath[1],
LocalPath: nameAndPath[1],
LoadMethod: loadMethod,
Debug: Settings.debug,
ParentName: parent,
ShardFunction: sfunc,
Partition: part,
TotalPartitions: total,
}
}
return configs
}
func ConfigsFromJsonUrl(url string) []*hfile.CollectionConfig {
log.Printf("[ConfigsFromJsonUrl] Fetching config from %s...\n", url)
res, err := http.Get(url)
if err != nil {
log.Fatal(err)
}
defer res.Body.Close()
var specs CollectionSpecList
if err := json.NewDecoder(res.Body).Decode(&specs); err != nil {
log.Fatal(err)
}
log.Printf("[ConfigsFromJsonUrl] Found %d collections.\n", len(specs.Collections))
ret := make([]*hfile.CollectionConfig, len(specs.Collections))
for i, spec := range specs.Collections {
if spec.Url != "" {
name := fmt.Sprintf("%s/%d", spec.Collection, spec.Partition)
loadMethod := hfile.CopiedToMem
if Settings.onDisk {
loadMethod = hfile.OnDisk
} else if Settings.mlock {
loadMethod = hfile.MemlockFile
}
if spec.Ondemand {
loadMethod = hfile.OnDisk
}
ret[i] = &hfile.CollectionConfig{
Name: name,
SourcePath: spec.Url,
LocalPath: "",
LoadMethod: loadMethod,
Debug: Settings.debug,
ParentName: spec.Collection,
ShardFunction: spec.Function,
Partition: fmt.Sprintf("%d", spec.Partition),
TotalPartitions: fmt.Sprintf("%d", spec.Capacity),
}
}
}
log.Printf("Found %d collections in config:", len(ret))
for _, cfg := range ret {
if Settings.debug {
log.Printf("\t%s (%s)", cfg.Name, cfg.SourcePath)
} else {
log.Printf("\t%s", cfg.Name)
}
}
return ret
}