diff --git a/client.go b/client.go index 04fa8299..4d90091a 100644 --- a/client.go +++ b/client.go @@ -129,6 +129,11 @@ func (c *Client) Do(r *Request) (*Response, error) { return nil, err } req.Header.Set("User-Agent", DefaultUserAgent) // Maybe https://codereview.appspot.com/7532043. + for name, values := range r.ExtraHeaders { + for _, value := range values { + req.Header.Add(name, value) + } + } resp, err := c.Doer.Do(req) if err != nil { diff --git a/cmd/metha-sync/main.go b/cmd/metha-sync/main.go index 58af7019..adbcaf36 100644 --- a/cmd/metha-sync/main.go +++ b/cmd/metha-sync/main.go @@ -4,12 +4,17 @@ import ( "flag" "fmt" "io/ioutil" + "net/http" "os" + "strings" "github.com/miku/metha" + "github.com/miku/metha/xflag" log "github.com/sirupsen/logrus" ) +var extraHeaders xflag.Array + func main() { format := flag.String("format", "oai_dc", "metadata format") @@ -31,6 +36,8 @@ func main() { logFile := flag.String("log", "", "filename to log to") logStderr := flag.Bool("log-errors-to-stderr", false, "Log errors and warnings to STDERR. If -log or -q are not given, write full log to STDOUT") + flag.Var(&extraHeaders, "H", `extra HTTP header to pass to requests (repeatable); e.g. -H "token: 123" `) + flag.Parse() if *version { @@ -83,6 +90,15 @@ func main() { log.AddHook(metha.NewCopyHook(os.Stderr)) } + var extra http.Header + for _, s := range extraHeaders { + parts := strings.SplitN(s, ":", 1) + if len(parts) != 2 { + log.Fatal(`extra headers notation is "Some-Key: Some-Value"`) + } + extra.Add(parts[0], parts[1]) + } + harvest, err := metha.NewHarvest(baseURL) if err != nil { log.Fatal(err) @@ -103,6 +119,7 @@ func main() { harvest.IgnoreHTTPErrors = *ignoreHTTPErrors harvest.SuppressFormatParameter = *suppressFormatParameter harvest.DailyInterval = *daily + harvest.ExtraHeaders = extra log.Printf("harvest: %+v", harvest) diff --git a/harvest.go b/harvest.go index d238634c..f6760231 100644 --- a/harvest.go +++ b/harvest.go @@ -7,6 +7,7 @@ import ( "fmt" "io/ioutil" "math/rand" + "net/http" "os" "os/signal" "path/filepath" @@ -63,6 +64,7 @@ type Harvest struct { MaxEmptyResponses int SuppressFormatParameter bool DailyInterval bool + ExtraHeaders http.Header // XXX: Lazy via sync.Once? Identify *Identify @@ -312,6 +314,7 @@ func (h *Harvest) runInterval(iv Interval) error { ResumptionToken: token, CleanBeforeDecode: h.CleanBeforeDecode, SuppressFormatParameter: h.SuppressFormatParameter, + ExtraHeaders: h.ExtraHeaders, } var filedate string @@ -416,7 +419,11 @@ func (h *Harvest) earliestDate() (time.Time, error) { // identify runs an OAI identify request and caches the result. func (h *Harvest) identify() error { - req := Request{Verb: "Identify", BaseURL: h.BaseURL} + req := Request{ + Verb: "Identify", + BaseURL: h.BaseURL, + ExtraHeaders: h.ExtraHeaders, + } c := CreateClient(30*time.Second, 2) diff --git a/request.go b/request.go index a2c82645..293e5494 100644 --- a/request.go +++ b/request.go @@ -4,6 +4,7 @@ import ( "bytes" "errors" "fmt" + "net/http" "net/url" "sort" "strings" @@ -30,6 +31,7 @@ type Request struct { ResumptionToken string CleanBeforeDecode bool SuppressFormatParameter bool + ExtraHeaders http.Header } // Values enhances the builtin url.Values. diff --git a/xflag/flag.go b/xflag/flag.go new file mode 100644 index 00000000..c4aa5ecd --- /dev/null +++ b/xflag/flag.go @@ -0,0 +1,25 @@ +// Package xflag add an additional flag type Array for repeated string flags. +// +// var f xflag.Array +// flag.Var(&f, "r", "some repeatable flag") +// +// flag.Parse() // $ command -r a -r b -r c +// for _, v := range f { ... } // []string{"a", "b", "c"} +// +package xflag + +import "strings" + +// ArrayFlags allows to store lists of flag values. +type Array []string + +// String representation. +func (f *Array) String() string { + return strings.Join(*f, ", ") +} + +// Set appends a value. +func (f *Array) Set(value string) error { + *f = append(*f, value) + return nil +}