This repository has been archived by the owner on Oct 3, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 5
/
pending_images.go
107 lines (92 loc) · 2.17 KB
/
pending_images.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
package db
import (
"database/sql"
"encoding/json"
"strings"
"github.com/bakape/boorufetch"
"github.com/bakape/captchouli/v2/common"
)
// Image data fetched from boorus pending processing by random selection
type PendingImage struct {
Rating boorufetch.Rating
MD5 [16]byte
TargetTag, URL string
Tags []string
}
// Check, if image is already on the list of pending images
func IsPendingImage(md5 [16]byte) (bool, error) {
return imageExists("pending_images", md5)
}
// Insert a new image pending processing
func InsertPendingImage(img PendingImage) (err error) {
lowercaseTags(img.Tags)
tags, err := json.Marshal(img.Tags)
if err != nil {
return
}
dbMu.Lock()
defer dbMu.Unlock()
_, err = sq.Insert("pending_images").
Columns("rating", "hash", "target_tag", "url", "tags").
Values(img.Rating, img.MD5[:], img.TargetTag, img.URL, tags).
Exec()
return
}
// Deletes random pending pending image for tag and returns it, if any
func PopRandomPendingImage(tag string) (img PendingImage, err error) {
tag = strings.ToLower(tag)
dbMu.Lock()
defer dbMu.Unlock()
err = InTransaction(func(tx *sql.Tx) (err error) {
var n int
err = sq.Select("count(*)").
From("pending_images").
Where("target_tag = ?", tag).
RunWith(tx).
QueryRow().
Scan(&n)
if err != nil {
return
}
if n == 0 {
return sql.ErrNoRows
}
var tags, md5 []byte
err = sq.Select("rating", "hash", "url", "tags").
From("pending_images").
Where("target_tag = ?", tag).
OrderBy("hash").
Offset(uint64(common.RandomInt(n))).
Limit(1).
RunWith(tx).
QueryRow().
Scan(&img.Rating, &md5, &img.URL, &tags)
if err != nil {
return
}
img.TargetTag = tag
copy(img.MD5[:], md5)
err = json.Unmarshal(tags, &img.Tags)
if err != nil {
return
}
_, err = sq.Delete("pending_images").
Where("hash = ?", img.MD5[:]).
RunWith(tx).
Exec()
return
})
return
}
// Count pending images for tag
func CountPending(tag string) (n int, err error) {
tag = strings.ToLower(tag)
dbMu.RLock()
defer dbMu.RUnlock()
err = sq.Select("count(*)").
From("pending_images").
Where("target_tag = ?", tag).
QueryRow().
Scan(&n)
return
}