From a05a02230e975314bb3a87e9f3a9e16371608880 Mon Sep 17 00:00:00 2001 From: Alex Denes Date: Sun, 27 Nov 2022 14:14:37 +0000 Subject: [PATCH] Generalise query functions and reuse request structs --- generic.go | 2 +- modules/e621/main.go | 64 ++++++++++++++++++++++------------------ modules/gelbooru/main.go | 44 ++++++++++++++------------- 3 files changed, 61 insertions(+), 49 deletions(-) diff --git a/generic.go b/generic.go index 99fda85..06ca081 100644 --- a/generic.go +++ b/generic.go @@ -27,7 +27,7 @@ import ( type Tags []string type Jobs uint -type QueryFunc (func(Tags, Jobs) ([]Media, error)) +type Page uint type Media struct { Source string diff --git a/modules/e621/main.go b/modules/e621/main.go index 62f3483..16900f4 100644 --- a/modules/e621/main.go +++ b/modules/e621/main.go @@ -34,6 +34,12 @@ import ( ) type ( + result struct { + media []Media + err error + pid Page + } + user_id uint64 post_id uint64 pool_id uint64 @@ -112,21 +118,32 @@ type ( } ) -type result struct { - media []Media - err error - pid uint -} - -func Query(tags Tags, j_max Jobs) (mr []Media, err error) { +func Query(uname string, api_key string, tags []string, j_max Jobs, tout time.Duration) (mr []Media, err error) { res_chan := make(chan result) var r_arr []result - for pid, rpid, ppid := uint(0), uint(0), uint(0); ; { - go run_job(tags, pid, res_chan) + q := &http.Request{ + URL: &url.URL{ + Scheme: "https", + Host: "e621.net", + Path: "/posts.json", + }, + Header: make(http.Header), + } + q.Header.Set("user-agent", "gomon/1.0 (caskd, https://git.redxen.eu/caskd/gomon)") + uq := q.URL.Query() + uq.Set("tags", strings.Join(tags, " ")) + if uname != "" && api_key != "" { + uq.Set("login", uname) + uq.Set("api_key", api_key) + } + q.URL.RawQuery = uq.Encode() + + for pid, rpid, ppid := Page(0), Page(0), Page(0); ; { + go run_job(q, tout, pid, res_chan) pid++ - if pid < uint(j_max) { + if pid < Page(j_max) { continue } @@ -161,12 +178,12 @@ func Query(tags Tags, j_max Jobs) (mr []Media, err error) { return } -func run_job(tags []string, pid uint, res chan result) { +func run_job(q *http.Request, tout time.Duration, pid Page, res chan result) { r := result{pid: pid} defer func(x *result, c chan result) { c <- *x }(&r, res) var rc io.ReadCloser - if rc, r.err = fetch(tags, pid); r.err != nil { + if rc, r.err = fetch(q, tout, pid); r.err != nil { return } defer rc.Close() @@ -176,23 +193,14 @@ func run_job(tags []string, pid uint, res chan result) { } } -func fetch(tags []string, pid uint) (rc io.ReadCloser, err error) { - client := http.Client{Timeout: 10 * time.Second} - req := &http.Request{ - URL: &url.URL{ - Scheme: "https", - Host: "e621.net", - Path: "/posts.json", - }, - } - query := req.URL.Query() - query.Add("page", strconv.FormatUint(uint64(pid), 10)) - query.Add("tags", strings.Join(tags, " ")) - req.URL.RawQuery = query.Encode() - req.Header = make(http.Header) - req.Header.Set("user-agent", "gomon/1.0 (https://git.redxen.eu/caskd/gomon)") +func fetch(q *http.Request, tout time.Duration, pid Page) (rc io.ReadCloser, err error) { + client := http.Client{Timeout: tout} - resp, err := client.Do(req) + uq := q.URL.Query() + uq.Set("page", strconv.FormatUint(uint64(pid), 10)) + q.URL.RawQuery = uq.Encode() + + resp, err := client.Do((*http.Request)(q)) if err != nil { return } diff --git a/modules/gelbooru/main.go b/modules/gelbooru/main.go index f498adb..8dce836 100644 --- a/modules/gelbooru/main.go +++ b/modules/gelbooru/main.go @@ -82,9 +82,24 @@ func Query(tags Tags, j_max Jobs) (mr []Media, err error) { res_chan := make(chan result) var r_arr []result + q := &http.Request{ + URL: &url.URL{ + Scheme: "https", + Host: "gelbooru.com", + Path: "/index.php", + }, + } + qu := q.URL.Query() + qu.Set("page", "dapi") + qu.Set("s", "post") + qu.Set("q", "index") + qu.Set("json", "1") + qu.Set("tags", strings.Join(tags, " ")) + q.URL.RawQuery = qu.Encode() + for pid, rpid, ppid := uint(0), uint(0), uint(0); ; { if pid <= 200 { // API only allows to fetch up to 200 pages per query - go run_job(tags, pid, res_chan) + go run_job(q, pid, res_chan) pid++ } @@ -123,12 +138,12 @@ func Query(tags Tags, j_max Jobs) (mr []Media, err error) { return } -func run_job(tags []string, pid uint, res chan result) { +func run_job(q *http.Request, pid uint, res chan result) { r := result{pid: pid} defer func(x *result, c chan result) { c <- *x }(&r, res) var rc io.ReadCloser - if rc, r.err = fetch(tags, pid); r.err != nil { + if rc, r.err = fetch(q, pid); r.err != nil { return } defer rc.Close() @@ -138,25 +153,14 @@ func run_job(tags []string, pid uint, res chan result) { } } -func fetch(tags []string, pid uint) (rc io.ReadCloser, err error) { +func fetch(q *http.Request, pid uint) (rc io.ReadCloser, err error) { client := http.Client{Timeout: 10 * time.Second} - req := &http.Request{ - URL: &url.URL{ - Scheme: "https", - Host: "gelbooru.com", - Path: "/index.php", - }, - } - query := req.URL.Query() - query.Add("page", "dapi") - query.Add("s", "post") - query.Add("q", "index") - query.Add("json", "1") - query.Add("pid", strconv.FormatUint(uint64(pid), 10)) - query.Add("tags", strings.Join(tags, " ")) - req.URL.RawQuery = query.Encode() - resp, err := client.Do(req) + qm := q.URL.Query() + qm.Set("pid", strconv.FormatUint(uint64(pid), 10)) + q.URL.RawQuery = qm.Encode() + + resp, err := client.Do(q) if err != nil { return }