├── .gitattributes ├── .github └── workflows │ └── releases.yml ├── .gitignore ├── LICENSE ├── README.md ├── common ├── convert │ ├── DBToRes.go │ └── RawToDB.go ├── database │ ├── drivers │ │ ├── cache.go │ │ ├── elasticsearch.go │ │ ├── filter.go │ │ ├── mongodb.go │ │ ├── neardb.go │ │ ├── pb │ │ │ ├── buildproto.bat │ │ │ ├── pb.pb.go │ │ │ ├── pb.proto │ │ │ └── pb_grpc.pb.go │ │ ├── rabbitmq.go │ │ └── redis.go │ ├── operations │ │ ├── illust.go │ │ ├── operations.go │ │ ├── rank.go │ │ ├── ugoira.go │ │ └── user.go │ └── tasktracer │ │ └── tracer.go ├── models │ ├── config.go │ ├── database.go │ ├── errors.go │ ├── interface.go │ ├── message.go │ ├── pixiv.go │ └── response.go └── utils │ ├── config │ └── const.go │ ├── hash.go │ ├── msgpack.go │ ├── string.go │ ├── telemetry │ ├── logger.go │ ├── loki.go │ └── prometheus.go │ └── typeconv.go ├── go.mod ├── go.sum ├── modules ├── responser │ ├── reader │ │ ├── dbreader.go │ │ ├── reader.go │ │ └── searchreader.go │ ├── responser.go │ ├── router.go │ └── task │ │ ├── gens.go │ │ └── task.go ├── spider │ ├── apis │ │ └── pixiv.go │ ├── pipeline │ │ ├── hook.go │ │ ├── pipeline.go │ │ └── response.go │ ├── scheduler │ │ ├── scheduler.go │ │ └── task.go │ ├── spider.go │ └── storage │ │ ├── storage.go │ │ └── utils.go ├── storer │ ├── handler.go │ ├── source │ │ ├── queue.go │ │ └── task.go │ └── storer.go └── util.go ├── responser.go ├── spider.go ├── storer.go ├── tasks.go └── tools └── incrementor └── incrementor.go /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.github/workflows/releases.yml: -------------------------------------------------------------------------------- 1 | name: Build&Release 2 | 3 | on: 4 | push: 5 | tags: 6 | - 'v*' 7 | - '*-rc' 8 | workflow_dispatch: 9 | 10 | jobs: 11 | build: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: actions/checkout@v2 15 | 16 | - name: Set up Go 17 | uses: actions/setup-go@v2 18 | with: 19 | go-version: 1.17 20 | 21 | - name: Build 22 | run: go build -o dist/ responser.go & go build -o dist/ spider.go & go build -o dist/ storer.go & go build -o dist/ tasks.go 23 | 24 | - name: Publish 25 | uses: ncipollo/release-action@v1 26 | with: 27 | artifacts: "./dist/*" 28 | commit: "main" 29 | token: ${{ secrets.GITHUB_TOKEN }} 30 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | .idea -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Shugetsu Soft 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Pixivel Backend 2 | The Backend Service for [pixivel.moe](https://pixivel.moe), written in golang. 3 | -------------------------------------------------------------------------------- /common/convert/DBToRes.go: -------------------------------------------------------------------------------- 1 | package convert 2 | 3 | import "github.com/ShugetsuSoft/pixivel-back/common/models" 4 | 5 | func Illust2IllustResponse(illust *models.Illust, user *models.User) *models.IllustResponse { 6 | var userres *models.UserResponse 7 | if user != nil { 8 | userres = User2UserResponse(user) 9 | } else { 10 | userres = nil 11 | } 12 | illustTags := Tags2TagResponses(illust.Tags) 13 | 14 | return &models.IllustResponse{ 15 | ID: illust.ID, 16 | Title: illust.Title, 17 | AltTitle: illust.AltTitle, 18 | Description: illust.Description, 19 | Type: illust.Type, 20 | CreateDate: illust.CreateDate.Format("2006-01-02T15:04:05"), 21 | UploadDate: illust.UploadDate.Format("2006-01-02T15:04:05"), 22 | Sanity: illust.Sanity, 23 | Width: illust.Width, 24 | Height: illust.Height, 25 | PageCount: illust.PageCount, 26 | Tags: illustTags, 27 | Statistic: models.IllustStatisticResponse{ 28 | Bookmarks: illust.Statistic.Bookmarks, 29 | Likes: illust.Statistic.Likes, 30 | Comments: illust.Statistic.Comments, 31 | Views: illust.Statistic.Views, 32 | }, 33 | User: userres, 34 | Image: illust.Image.Format("2006-01-02T15:04:05"), 35 | AIType: illust.AIType, 36 | } 37 | } 38 | 39 | func User2UserResponse(user *models.User) *models.UserResponse { 40 | return &models.UserResponse{ 41 | ID: user.ID, 42 | Name: user.Name, 43 | Bio: user.Bio, 44 | Image: models.UserImageResponse{ 45 | Url: user.Image.Url, 46 | BigUrl: user.Image.BigUrl, 47 | Background: user.Image.Background, 48 | }, 49 | } 50 | } 51 | 52 | func Illusts2IllustsResponse(illusts []models.Illust, hasnext bool) *models.IllustsResponse { 53 | illustsresponse := make([]models.IllustResponse, len(illusts)) 54 | for i := range illusts { 55 | illustsresponse[i] = *Illust2IllustResponse(&illusts[i], nil) 56 | } 57 | return &models.IllustsResponse{ 58 | Illusts: illustsresponse, 59 | HasNext: hasnext, 60 | } 61 | } 62 | 63 | func Illusts2IllustsSearchResponse(illusts []models.Illust, hasnext bool, scores []float64, highlights []*string) *models.IllustsSearchResponse { 64 | illustsresponse := make([]models.IllustResponse, len(illusts)) 65 | for i := range illusts { 66 | illustsresponse[i] = *Illust2IllustResponse(&illusts[i], nil) 67 | } 68 | return &models.IllustsSearchResponse{ 69 | Illusts: illustsresponse, 70 | Scores: scores, 71 | HighLight: highlights, 72 | HasNext: hasnext, 73 | } 74 | } 75 | 76 | func Users2UsersSearchResponse(users []models.User, hasnext bool, scores []float64, highlights []*string) *models.UsersSearchResponse { 77 | usersesponse := make([]models.UserResponse, len(users)) 78 | for i := range users { 79 | usersesponse[i] = *User2UserResponse(&users[i]) 80 | } 81 | return &models.UsersSearchResponse{ 82 | Users: usersesponse, 83 | Scores: scores, 84 | HighLight: highlights, 85 | HasNext: hasnext, 86 | } 87 | } 88 | 89 | func Tags2TagResponses(tags []models.IllustTag) []models.IllustTagResponse { 90 | var illustTags []models.IllustTagResponse 91 | 92 | illustTags = make([]models.IllustTagResponse, len(tags)) 93 | for i := 0; i < len(tags); i++ { 94 | illustTags[i] = models.IllustTagResponse{ 95 | Name: tags[i].Name, 96 | Translation: tags[i].Translation, 97 | } 98 | } 99 | return illustTags 100 | } 101 | 102 | func RankAggregateResult2IllustsResponses(rank []models.RankAggregateResult, hasnext bool) *models.IllustsResponse { 103 | illusts := make([]models.Illust, len(rank)) 104 | for i := range rank { 105 | illusts[i] = rank[i].Content 106 | } 107 | return Illusts2IllustsResponse(illusts, hasnext) 108 | } 109 | 110 | func Ugoira2UgoiraResponse(ugoira *models.Ugoira) *models.UgoiraResponse { 111 | frames := make([]models.UgoiraFrameResponse, len(ugoira.Frames)) 112 | for key := range ugoira.Frames { 113 | frames[key].Delay = ugoira.Frames[key].Delay 114 | frames[key].File = ugoira.Frames[key].File 115 | } 116 | return &models.UgoiraResponse{ 117 | Image: ugoira.Image.Format("2006-01-02T15:04:05"), 118 | MimeType: ugoira.MimeType, 119 | Frames: frames, 120 | ID: ugoira.ID, 121 | } 122 | } 123 | 124 | func Users2UsersResponse(users []models.User, hasnext bool) *models.UsersResponse { 125 | usersresponse := make([]models.UserResponse, len(users)) 126 | for i := range users { 127 | usersresponse[i] = *User2UserResponse(&users[i]) 128 | } 129 | return &models.UsersResponse{ 130 | Users: usersresponse, 131 | HasNext: hasnext, 132 | } 133 | } 134 | -------------------------------------------------------------------------------- /common/convert/RawToDB.go: -------------------------------------------------------------------------------- 1 | package convert 2 | 3 | import ( 4 | "strings" 5 | "time" 6 | 7 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 8 | 9 | "github.com/ShugetsuSoft/pixivel-back/common/models" 10 | ) 11 | 12 | var banTagList = map[string]bool{ 13 | "+18": true, 14 | "nude": true, 15 | "r18": true, 16 | "r-18": true, 17 | "r-18g": true, 18 | } 19 | 20 | func ParseImgTime(url string) time.Time { 21 | date_str := url[strings.Index(url, "/img/")+5 : strings.LastIndex(url, "/")] 22 | t, _ := time.Parse("2006/01/02/15/04/05", date_str) 23 | return t 24 | } 25 | 26 | func IsIllustBanned(raw *models.IllustRaw) bool { 27 | if raw.XRestrict == 1 { 28 | return true 29 | } 30 | for i := 0; i < len(raw.Tags.Tags); i++ { 31 | if _, is := banTagList[strings.ToLower(raw.Tags.Tags[i].Tag)]; is { 32 | return true 33 | } 34 | } 35 | return false 36 | } 37 | 38 | func IllustRaw2Illust(raw *models.IllustRaw) *models.Illust { 39 | illustTags := make([]models.IllustTag, len(raw.Tags.Tags)) 40 | for i := 0; i < len(raw.Tags.Tags); i++ { 41 | illustTags[i] = models.IllustTag{ 42 | Name: raw.Tags.Tags[i].Tag, 43 | Translation: raw.Tags.Tags[i].Translation.En, 44 | } 45 | } 46 | sta := models.IllustStatistic{ 47 | Bookmarks: raw.BookmarkCount, 48 | Likes: raw.LikeCount, 49 | Comments: raw.CommentCount, 50 | Views: raw.ViewCount, 51 | } 52 | return &models.Illust{ 53 | ID: raw.ID, 54 | Title: raw.Title, 55 | AltTitle: raw.Alt, 56 | Description: raw.Description, 57 | Type: raw.IllustType, 58 | CreateDate: raw.CreateDate, 59 | UploadDate: raw.UploadDate, 60 | Sanity: raw.Sl, 61 | Width: raw.Width, 62 | Height: raw.Height, 63 | PageCount: raw.PageCount, 64 | User: raw.UserID, 65 | Tags: illustTags, 66 | Popularity: CalcIllustPop(sta), 67 | Statistic: sta, 68 | Image: ParseImgTime(raw.Urls.Original), 69 | Banned: IsIllustBanned(raw), 70 | AIType: raw.AIType, 71 | } 72 | } 73 | 74 | func CalcIllustPop(sta models.IllustStatistic) uint { 75 | return (sta.Bookmarks*70 + sta.Likes*30) / 100 76 | } 77 | 78 | func Illust2IllustSearch(illustdb *models.Illust) *models.IllustSearch { 79 | return &models.IllustSearch{ 80 | Title: utils.RemoveSpecialChars(illustdb.Title), 81 | AltTitle: utils.RemoveSpecialChars(illustdb.AltTitle), 82 | Description: utils.RemoveSpecialChars(illustdb.Description), 83 | Type: illustdb.Type, 84 | CreateDate: illustdb.CreateDate, 85 | Sanity: illustdb.Sanity, 86 | Popularity: illustdb.Popularity, 87 | User: illustdb.User, 88 | Tags: utils.RemoveSpecialCharsTags(illustdb.Tags), 89 | } 90 | } 91 | 92 | func User2UserSearch(userdb *models.User) *models.UserSearch { 93 | return &models.UserSearch{ 94 | Name: utils.RemoveSpecialChars(userdb.Name), 95 | Bio: utils.RemoveSpecialChars(userdb.Bio), 96 | } 97 | } 98 | 99 | func UserRaw2User(raw *models.UserRaw) *models.User { 100 | return &models.User{ 101 | ID: raw.UserID, 102 | Name: raw.Name, 103 | Bio: raw.Comment, 104 | Image: models.UserImage{ 105 | Url: raw.Image, 106 | BigUrl: raw.ImageBig, 107 | Background: raw.Background.URL, 108 | }, 109 | } 110 | } 111 | 112 | func UserIllusts2UserIllustsResponse(uid uint64, raw *models.UserIllustsRaw) *models.UserIllustsResponse { 113 | lenth := 0 114 | switch illusts := raw.Illusts.(type) { 115 | case map[string]interface{}: 116 | lenth += len(illusts) 117 | } 118 | switch manga := raw.Manga.(type) { 119 | case map[string]interface{}: 120 | lenth += len(manga) 121 | } 122 | lis := make(models.IDList, lenth) 123 | i := 0 124 | switch illusts := raw.Illusts.(type) { 125 | case map[string]interface{}: 126 | for key := range illusts { 127 | lis[i] = utils.Atoi(key) 128 | i++ 129 | } 130 | } 131 | switch manga := raw.Manga.(type) { 132 | case map[string]interface{}: 133 | for key := range manga { 134 | lis[i] = utils.Atoi(key) 135 | i++ 136 | } 137 | } 138 | return &models.UserIllustsResponse{ 139 | UserID: uid, 140 | Illusts: lis, 141 | } 142 | } 143 | 144 | func RankIllusts2RankIllustsResponse(raw *models.RankIllustsRawResponse) *models.RankIllustsResponseMessage { 145 | illusts := make([]models.IDWithPos, len(raw.Contents)) 146 | for key := range raw.Contents { 147 | illusts[key].ID = raw.Contents[key].IllustID 148 | illusts[key].Pos = raw.Contents[key].Rank 149 | } 150 | hasnext := true 151 | switch next := raw.Next.(type) { 152 | case bool: 153 | if next == false { 154 | hasnext = false 155 | } 156 | } 157 | return &models.RankIllustsResponseMessage{ 158 | Mode: raw.Mode, 159 | Date: raw.Date, 160 | Content: raw.Content, 161 | Page: raw.Page, 162 | Next: hasnext, 163 | Illusts: illusts, 164 | } 165 | } 166 | 167 | func UgoiraRaw2Ugoira(raw *models.UgoiraRaw, id uint64) *models.Ugoira { 168 | frames := make([]models.UgoiraFrame, len(raw.Frames)) 169 | for key := range raw.Frames { 170 | frames[key].Delay = raw.Frames[key].Delay 171 | frames[key].File = raw.Frames[key].File 172 | } 173 | return &models.Ugoira{ 174 | Image: ParseImgTime(raw.Src), 175 | MimeType: raw.MimeType, 176 | Frames: frames, 177 | ID: id, 178 | } 179 | } 180 | -------------------------------------------------------------------------------- /common/database/drivers/cache.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import ( 4 | "bytes" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 6 | "github.com/gomodule/redigo/redis" 7 | "github.com/vmihailenco/msgpack/v5" 8 | ) 9 | 10 | type Cache struct { 11 | redis *RedisPool 12 | } 13 | 14 | func NewCache(redis *RedisPool) *Cache { 15 | return &Cache{ 16 | redis: redis, 17 | } 18 | } 19 | 20 | func marshal(v interface{}) ([]byte, error) { 21 | enc := msgpack.NewEncoder(nil) 22 | var buf bytes.Buffer 23 | enc.Reset(&buf) 24 | enc.SetCustomStructTag("json") 25 | 26 | err := enc.Encode(v) 27 | b := buf.Bytes() 28 | 29 | if err != nil { 30 | return nil, err 31 | } 32 | return b, err 33 | } 34 | 35 | func unmarshal(data []byte, v interface{}) error { 36 | dec := msgpack.NewDecoder(nil) 37 | dec.Reset(bytes.NewReader(data)) 38 | dec.SetCustomStructTag("json") 39 | 40 | err := dec.Decode(v) 41 | 42 | return err 43 | } 44 | 45 | func calcKey(api string, params []string) string { 46 | key := api + "-" 47 | for k := range params { 48 | key += params[k] + "." 49 | } 50 | return key 51 | } 52 | 53 | func (c *Cache) Get(api string, params ...string) (interface{}, error) { 54 | key := calcKey(api, params) 55 | cli := c.redis.NewRedisClient() 56 | defer cli.Close() 57 | cached, err := cli.GetValueBytes(key) 58 | if err != nil { 59 | if err == redis.ErrNil { 60 | return nil, nil 61 | } 62 | return nil, err 63 | } 64 | var result interface{} 65 | err = unmarshal(cached, &result) 66 | telemetry.RequestsHitCache.Inc() 67 | return result, err 68 | } 69 | 70 | func (c *Cache) Set(api string, value interface{}, expire int, params ...string) error { 71 | key := calcKey(api, params) 72 | cli := c.redis.NewRedisClient() 73 | defer cli.Close() 74 | binval, err := marshal(value) 75 | if err != nil { 76 | return err 77 | } 78 | return cli.SetValueBytesExpire(key, binval, expire) 79 | } 80 | 81 | func (c *Cache) Clear(api string, params ...string) error { 82 | key := calcKey(api, params) 83 | cli := c.redis.NewRedisClient() 84 | defer cli.Close() 85 | err := cli.DeleteValue(key) 86 | if err == redis.ErrNil { 87 | return nil 88 | } 89 | return err 90 | } 91 | -------------------------------------------------------------------------------- /common/database/drivers/elasticsearch.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/ShugetsuSoft/pixivel-back/common/models" 7 | elastic "github.com/olivere/elastic/v7" 8 | ) 9 | 10 | type ElasticSearch struct { 11 | cli *elastic.Client 12 | ctx context.Context 13 | } 14 | 15 | func NewElasticSearchClient(ctx context.Context, uri string, username string, password string) (*ElasticSearch, error) { 16 | client, err := elastic.NewClient(elastic.SetURL(uri), elastic.SetBasicAuth(username, password), elastic.SetSniff(false)) 17 | if err != nil { 18 | return nil, err 19 | } 20 | return &ElasticSearch{ 21 | cli: client, 22 | ctx: ctx, 23 | }, nil 24 | } 25 | 26 | func (es *ElasticSearch) CreateIndex(ctx context.Context, name string, mapping string) error { 27 | exists, err := es.cli.IndexExists(name).Do(ctx) 28 | if err != nil { 29 | return err 30 | } 31 | if !exists { 32 | _, err := es.cli.CreateIndex(name).BodyString(mapping).Do(ctx) 33 | if err != nil { 34 | return err 35 | } 36 | } else { 37 | return models.ErrorIndexExist 38 | } 39 | return nil 40 | } 41 | 42 | func (es *ElasticSearch) DeleteIndex(ctx context.Context, name string) error { 43 | _, err := es.cli.DeleteIndex(name).Do(ctx) 44 | if err != nil { 45 | return err 46 | } 47 | return nil 48 | } 49 | 50 | func (es *ElasticSearch) InsertDocument(ctx context.Context, index string, id string, object interface{}) error { 51 | _, err := es.cli.Index(). 52 | Index(index). 53 | Id(id). 54 | BodyJson(object). 55 | OpType("index"). 56 | Do(ctx) 57 | if err != nil { 58 | return err 59 | } 60 | return nil 61 | } 62 | 63 | func (es *ElasticSearch) Search(index string) *elastic.SearchService { 64 | return es.cli.Search().Index(index) 65 | } 66 | 67 | func (es *ElasticSearch) MultiSearch() *elastic.MultiSearchService { 68 | return es.cli.MultiSearch() 69 | } 70 | 71 | func (es *ElasticSearch) Query(key string, val interface{}) *elastic.MatchQuery { 72 | return elastic.NewMatchQuery(key, val) 73 | } 74 | 75 | func (es *ElasticSearch) TermsQuery(key string, val []string) *elastic.TermsQuery { 76 | return elastic.NewTermsQueryFromStrings(key, val...) 77 | } 78 | 79 | func (es *ElasticSearch) Suggest(key string) *elastic.CompletionSuggester { 80 | return elastic.NewCompletionSuggester(key) 81 | } 82 | 83 | func (es *ElasticSearch) TermSuggest(key string) *elastic.TermSuggester { 84 | return elastic.NewTermSuggester(key) 85 | } 86 | 87 | func (es *ElasticSearch) BoolQuery() *elastic.BoolQuery { 88 | return elastic.NewBoolQuery() 89 | } 90 | 91 | func (es *ElasticSearch) DoSearch(ctx context.Context, s *elastic.SearchService) (*elastic.SearchResult, error) { 92 | return s.Do(ctx) 93 | } 94 | 95 | func (es *ElasticSearch) DoMultiSearch(ctx context.Context, s *elastic.MultiSearchService) (*elastic.MultiSearchResult, error) { 96 | return s.Do(ctx) 97 | } 98 | -------------------------------------------------------------------------------- /common/database/drivers/filter.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import redisbloom "github.com/RedisBloom/redisbloom-go" 4 | 5 | type BloomFilter struct { 6 | cli *redisbloom.Client 7 | } 8 | 9 | func (red *RedisPool) NewBloomFilter(name string) *BloomFilter { 10 | client := redisbloom.NewClientFromPool(red.RedisPool, name) 11 | return &BloomFilter{ 12 | cli: client, 13 | } 14 | } 15 | 16 | func NewBloomFilter(pool *RedisPool, name string) *BloomFilter { 17 | client := redisbloom.NewClientFromPool(pool.RedisPool, name) 18 | return &BloomFilter{ 19 | cli: client, 20 | } 21 | } 22 | 23 | func (ckft *BloomFilter) Add(name string, key string) (bool, error) { 24 | return ckft.cli.Add(name, key) 25 | } 26 | 27 | func (ckft *BloomFilter) Exists(name string, key string) (bool, error) { 28 | return ckft.cli.Exists(name, key) 29 | } 30 | -------------------------------------------------------------------------------- /common/database/drivers/mongodb.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import ( 4 | "context" 5 | 6 | "go.mongodb.org/mongo-driver/mongo" 7 | "go.mongodb.org/mongo-driver/mongo/options" 8 | ) 9 | 10 | type MongoDatabase struct { 11 | Context context.Context 12 | Client *mongo.Client 13 | DB *mongo.Database 14 | } 15 | 16 | func NewMongoDatabase(context context.Context, uri string, dbname string) (*MongoDatabase, error) { 17 | clientOptions := options.Client().ApplyURI(uri).SetMaxPoolSize(20) 18 | client, err := mongo.Connect(context, clientOptions) 19 | if err != nil { 20 | return nil, err 21 | } 22 | 23 | err = client.Ping(context, nil) 24 | if err != nil { 25 | return nil, err 26 | } 27 | 28 | return &MongoDatabase{ 29 | Context: context, 30 | Client: client, 31 | DB: client.Database(dbname), 32 | }, nil 33 | } 34 | 35 | func (db *MongoDatabase) Session() (mongo.Session, error) { 36 | session, err := db.Client.StartSession() 37 | if err != nil { 38 | return nil, err 39 | } 40 | return session, nil 41 | } 42 | 43 | func (db *MongoDatabase) Collection(colname string) *mongo.Collection { 44 | return db.DB.Collection(colname) 45 | } 46 | 47 | func (db *MongoDatabase) Close() error { 48 | return db.Client.Disconnect(db.Context) 49 | } 50 | -------------------------------------------------------------------------------- /common/database/drivers/neardb.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/ShugetsuSoft/pixivel-back/common/database/drivers/pb" 7 | "google.golang.org/grpc" 8 | ) 9 | 10 | type NearDB struct { 11 | Conn *grpc.ClientConn 12 | Client pb.NearDBServiceClient 13 | } 14 | 15 | func NewNearDB(uri string) (*NearDB, error) { 16 | conn, err := grpc.Dial(uri, grpc.WithInsecure()) 17 | if err != nil { 18 | return nil, err 19 | } 20 | client := pb.NewNearDBServiceClient(conn) 21 | return &NearDB{ 22 | Conn: conn, 23 | Client: client, 24 | }, nil 25 | } 26 | 27 | func (near *NearDB) Add(ctx context.Context, id uint64, set []string, pop uint) error { 28 | _, err := near.Client.Add(ctx, &pb.AddRequest{ 29 | Id: id, 30 | Taglist: set, 31 | Pop: uint64(pop), 32 | }) 33 | return err 34 | } 35 | 36 | func (near *NearDB) Query(ctx context.Context, set []string, k int, drif float64) ([]*pb.Item, error) { 37 | items, err := near.Client.Query(ctx, &pb.QueryRequest{ 38 | Taglist: set, 39 | K: int64(k), 40 | Drift: drif, 41 | }) 42 | return items.GetItems(), err 43 | } 44 | 45 | func (near *NearDB) QueryById(ctx context.Context, id uint64, k int, drif float64) ([]*pb.Item, error) { 46 | items, err := near.Client.QueryById(ctx, &pb.QueryByIdRequest{ 47 | Id: id, 48 | K: int64(k), 49 | Drift: drif, 50 | }) 51 | return items.GetItems(), err 52 | } 53 | -------------------------------------------------------------------------------- /common/database/drivers/pb/buildproto.bat: -------------------------------------------------------------------------------- 1 | @echo off 2 | protoc --go-grpc_out=. --go_out=. ./*.proto 3 | pause -------------------------------------------------------------------------------- /common/database/drivers/pb/pb.pb.go: -------------------------------------------------------------------------------- 1 | // Code generated by protoc-gen-go. DO NOT EDIT. 2 | // versions: 3 | // protoc-gen-go v1.28.0 4 | // protoc v3.21.2 5 | // source: pb.proto 6 | 7 | package pb 8 | 9 | import ( 10 | protoreflect "google.golang.org/protobuf/reflect/protoreflect" 11 | protoimpl "google.golang.org/protobuf/runtime/protoimpl" 12 | reflect "reflect" 13 | sync "sync" 14 | ) 15 | 16 | const ( 17 | // Verify that this generated code is sufficiently up-to-date. 18 | _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) 19 | // Verify that runtime/protoimpl is sufficiently up-to-date. 20 | _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) 21 | ) 22 | 23 | type AddRequest struct { 24 | state protoimpl.MessageState 25 | sizeCache protoimpl.SizeCache 26 | unknownFields protoimpl.UnknownFields 27 | 28 | Id uint64 `protobuf:"varint,1,opt,name=id,proto3" json:"id,omitempty"` 29 | Taglist []string `protobuf:"bytes,2,rep,name=taglist,proto3" json:"taglist,omitempty"` 30 | Pop uint64 `protobuf:"varint,3,opt,name=pop,proto3" json:"pop,omitempty"` 31 | } 32 | 33 | func (x *AddRequest) Reset() { 34 | *x = AddRequest{} 35 | if protoimpl.UnsafeEnabled { 36 | mi := &file_pb_proto_msgTypes[0] 37 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 38 | ms.StoreMessageInfo(mi) 39 | } 40 | } 41 | 42 | func (x *AddRequest) String() string { 43 | return protoimpl.X.MessageStringOf(x) 44 | } 45 | 46 | func (*AddRequest) ProtoMessage() {} 47 | 48 | func (x *AddRequest) ProtoReflect() protoreflect.Message { 49 | mi := &file_pb_proto_msgTypes[0] 50 | if protoimpl.UnsafeEnabled && x != nil { 51 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 52 | if ms.LoadMessageInfo() == nil { 53 | ms.StoreMessageInfo(mi) 54 | } 55 | return ms 56 | } 57 | return mi.MessageOf(x) 58 | } 59 | 60 | // Deprecated: Use AddRequest.ProtoReflect.Descriptor instead. 61 | func (*AddRequest) Descriptor() ([]byte, []int) { 62 | return file_pb_proto_rawDescGZIP(), []int{0} 63 | } 64 | 65 | func (x *AddRequest) GetId() uint64 { 66 | if x != nil { 67 | return x.Id 68 | } 69 | return 0 70 | } 71 | 72 | func (x *AddRequest) GetTaglist() []string { 73 | if x != nil { 74 | return x.Taglist 75 | } 76 | return nil 77 | } 78 | 79 | func (x *AddRequest) GetPop() uint64 { 80 | if x != nil { 81 | return x.Pop 82 | } 83 | return 0 84 | } 85 | 86 | type RemoveRequest struct { 87 | state protoimpl.MessageState 88 | sizeCache protoimpl.SizeCache 89 | unknownFields protoimpl.UnknownFields 90 | 91 | Id uint64 `protobuf:"varint,4,opt,name=id,proto3" json:"id,omitempty"` 92 | } 93 | 94 | func (x *RemoveRequest) Reset() { 95 | *x = RemoveRequest{} 96 | if protoimpl.UnsafeEnabled { 97 | mi := &file_pb_proto_msgTypes[1] 98 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 99 | ms.StoreMessageInfo(mi) 100 | } 101 | } 102 | 103 | func (x *RemoveRequest) String() string { 104 | return protoimpl.X.MessageStringOf(x) 105 | } 106 | 107 | func (*RemoveRequest) ProtoMessage() {} 108 | 109 | func (x *RemoveRequest) ProtoReflect() protoreflect.Message { 110 | mi := &file_pb_proto_msgTypes[1] 111 | if protoimpl.UnsafeEnabled && x != nil { 112 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 113 | if ms.LoadMessageInfo() == nil { 114 | ms.StoreMessageInfo(mi) 115 | } 116 | return ms 117 | } 118 | return mi.MessageOf(x) 119 | } 120 | 121 | // Deprecated: Use RemoveRequest.ProtoReflect.Descriptor instead. 122 | func (*RemoveRequest) Descriptor() ([]byte, []int) { 123 | return file_pb_proto_rawDescGZIP(), []int{1} 124 | } 125 | 126 | func (x *RemoveRequest) GetId() uint64 { 127 | if x != nil { 128 | return x.Id 129 | } 130 | return 0 131 | } 132 | 133 | type NoneResponse struct { 134 | state protoimpl.MessageState 135 | sizeCache protoimpl.SizeCache 136 | unknownFields protoimpl.UnknownFields 137 | } 138 | 139 | func (x *NoneResponse) Reset() { 140 | *x = NoneResponse{} 141 | if protoimpl.UnsafeEnabled { 142 | mi := &file_pb_proto_msgTypes[2] 143 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 144 | ms.StoreMessageInfo(mi) 145 | } 146 | } 147 | 148 | func (x *NoneResponse) String() string { 149 | return protoimpl.X.MessageStringOf(x) 150 | } 151 | 152 | func (*NoneResponse) ProtoMessage() {} 153 | 154 | func (x *NoneResponse) ProtoReflect() protoreflect.Message { 155 | mi := &file_pb_proto_msgTypes[2] 156 | if protoimpl.UnsafeEnabled && x != nil { 157 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 158 | if ms.LoadMessageInfo() == nil { 159 | ms.StoreMessageInfo(mi) 160 | } 161 | return ms 162 | } 163 | return mi.MessageOf(x) 164 | } 165 | 166 | // Deprecated: Use NoneResponse.ProtoReflect.Descriptor instead. 167 | func (*NoneResponse) Descriptor() ([]byte, []int) { 168 | return file_pb_proto_rawDescGZIP(), []int{2} 169 | } 170 | 171 | type QueryRequest struct { 172 | state protoimpl.MessageState 173 | sizeCache protoimpl.SizeCache 174 | unknownFields protoimpl.UnknownFields 175 | 176 | Taglist []string `protobuf:"bytes,5,rep,name=taglist,proto3" json:"taglist,omitempty"` 177 | K int64 `protobuf:"varint,6,opt,name=k,proto3" json:"k,omitempty"` 178 | Drift float64 `protobuf:"fixed64,7,opt,name=drift,proto3" json:"drift,omitempty"` 179 | } 180 | 181 | func (x *QueryRequest) Reset() { 182 | *x = QueryRequest{} 183 | if protoimpl.UnsafeEnabled { 184 | mi := &file_pb_proto_msgTypes[3] 185 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 186 | ms.StoreMessageInfo(mi) 187 | } 188 | } 189 | 190 | func (x *QueryRequest) String() string { 191 | return protoimpl.X.MessageStringOf(x) 192 | } 193 | 194 | func (*QueryRequest) ProtoMessage() {} 195 | 196 | func (x *QueryRequest) ProtoReflect() protoreflect.Message { 197 | mi := &file_pb_proto_msgTypes[3] 198 | if protoimpl.UnsafeEnabled && x != nil { 199 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 200 | if ms.LoadMessageInfo() == nil { 201 | ms.StoreMessageInfo(mi) 202 | } 203 | return ms 204 | } 205 | return mi.MessageOf(x) 206 | } 207 | 208 | // Deprecated: Use QueryRequest.ProtoReflect.Descriptor instead. 209 | func (*QueryRequest) Descriptor() ([]byte, []int) { 210 | return file_pb_proto_rawDescGZIP(), []int{3} 211 | } 212 | 213 | func (x *QueryRequest) GetTaglist() []string { 214 | if x != nil { 215 | return x.Taglist 216 | } 217 | return nil 218 | } 219 | 220 | func (x *QueryRequest) GetK() int64 { 221 | if x != nil { 222 | return x.K 223 | } 224 | return 0 225 | } 226 | 227 | func (x *QueryRequest) GetDrift() float64 { 228 | if x != nil { 229 | return x.Drift 230 | } 231 | return 0 232 | } 233 | 234 | type Item struct { 235 | state protoimpl.MessageState 236 | sizeCache protoimpl.SizeCache 237 | unknownFields protoimpl.UnknownFields 238 | 239 | Id uint64 `protobuf:"varint,8,opt,name=id,proto3" json:"id,omitempty"` 240 | Distance float32 `protobuf:"fixed32,9,opt,name=distance,proto3" json:"distance,omitempty"` 241 | } 242 | 243 | func (x *Item) Reset() { 244 | *x = Item{} 245 | if protoimpl.UnsafeEnabled { 246 | mi := &file_pb_proto_msgTypes[4] 247 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 248 | ms.StoreMessageInfo(mi) 249 | } 250 | } 251 | 252 | func (x *Item) String() string { 253 | return protoimpl.X.MessageStringOf(x) 254 | } 255 | 256 | func (*Item) ProtoMessage() {} 257 | 258 | func (x *Item) ProtoReflect() protoreflect.Message { 259 | mi := &file_pb_proto_msgTypes[4] 260 | if protoimpl.UnsafeEnabled && x != nil { 261 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 262 | if ms.LoadMessageInfo() == nil { 263 | ms.StoreMessageInfo(mi) 264 | } 265 | return ms 266 | } 267 | return mi.MessageOf(x) 268 | } 269 | 270 | // Deprecated: Use Item.ProtoReflect.Descriptor instead. 271 | func (*Item) Descriptor() ([]byte, []int) { 272 | return file_pb_proto_rawDescGZIP(), []int{4} 273 | } 274 | 275 | func (x *Item) GetId() uint64 { 276 | if x != nil { 277 | return x.Id 278 | } 279 | return 0 280 | } 281 | 282 | func (x *Item) GetDistance() float32 { 283 | if x != nil { 284 | return x.Distance 285 | } 286 | return 0 287 | } 288 | 289 | type QueryResponse struct { 290 | state protoimpl.MessageState 291 | sizeCache protoimpl.SizeCache 292 | unknownFields protoimpl.UnknownFields 293 | 294 | Items []*Item `protobuf:"bytes,10,rep,name=items,proto3" json:"items,omitempty"` 295 | } 296 | 297 | func (x *QueryResponse) Reset() { 298 | *x = QueryResponse{} 299 | if protoimpl.UnsafeEnabled { 300 | mi := &file_pb_proto_msgTypes[5] 301 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 302 | ms.StoreMessageInfo(mi) 303 | } 304 | } 305 | 306 | func (x *QueryResponse) String() string { 307 | return protoimpl.X.MessageStringOf(x) 308 | } 309 | 310 | func (*QueryResponse) ProtoMessage() {} 311 | 312 | func (x *QueryResponse) ProtoReflect() protoreflect.Message { 313 | mi := &file_pb_proto_msgTypes[5] 314 | if protoimpl.UnsafeEnabled && x != nil { 315 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 316 | if ms.LoadMessageInfo() == nil { 317 | ms.StoreMessageInfo(mi) 318 | } 319 | return ms 320 | } 321 | return mi.MessageOf(x) 322 | } 323 | 324 | // Deprecated: Use QueryResponse.ProtoReflect.Descriptor instead. 325 | func (*QueryResponse) Descriptor() ([]byte, []int) { 326 | return file_pb_proto_rawDescGZIP(), []int{5} 327 | } 328 | 329 | func (x *QueryResponse) GetItems() []*Item { 330 | if x != nil { 331 | return x.Items 332 | } 333 | return nil 334 | } 335 | 336 | type QueryByIdRequest struct { 337 | state protoimpl.MessageState 338 | sizeCache protoimpl.SizeCache 339 | unknownFields protoimpl.UnknownFields 340 | 341 | Id uint64 `protobuf:"varint,11,opt,name=id,proto3" json:"id,omitempty"` 342 | K int64 `protobuf:"varint,12,opt,name=k,proto3" json:"k,omitempty"` 343 | Drift float64 `protobuf:"fixed64,13,opt,name=drift,proto3" json:"drift,omitempty"` 344 | } 345 | 346 | func (x *QueryByIdRequest) Reset() { 347 | *x = QueryByIdRequest{} 348 | if protoimpl.UnsafeEnabled { 349 | mi := &file_pb_proto_msgTypes[6] 350 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 351 | ms.StoreMessageInfo(mi) 352 | } 353 | } 354 | 355 | func (x *QueryByIdRequest) String() string { 356 | return protoimpl.X.MessageStringOf(x) 357 | } 358 | 359 | func (*QueryByIdRequest) ProtoMessage() {} 360 | 361 | func (x *QueryByIdRequest) ProtoReflect() protoreflect.Message { 362 | mi := &file_pb_proto_msgTypes[6] 363 | if protoimpl.UnsafeEnabled && x != nil { 364 | ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) 365 | if ms.LoadMessageInfo() == nil { 366 | ms.StoreMessageInfo(mi) 367 | } 368 | return ms 369 | } 370 | return mi.MessageOf(x) 371 | } 372 | 373 | // Deprecated: Use QueryByIdRequest.ProtoReflect.Descriptor instead. 374 | func (*QueryByIdRequest) Descriptor() ([]byte, []int) { 375 | return file_pb_proto_rawDescGZIP(), []int{6} 376 | } 377 | 378 | func (x *QueryByIdRequest) GetId() uint64 { 379 | if x != nil { 380 | return x.Id 381 | } 382 | return 0 383 | } 384 | 385 | func (x *QueryByIdRequest) GetK() int64 { 386 | if x != nil { 387 | return x.K 388 | } 389 | return 0 390 | } 391 | 392 | func (x *QueryByIdRequest) GetDrift() float64 { 393 | if x != nil { 394 | return x.Drift 395 | } 396 | return 0 397 | } 398 | 399 | var File_pb_proto protoreflect.FileDescriptor 400 | 401 | var file_pb_proto_rawDesc = []byte{ 402 | 0x0a, 0x08, 0x70, 0x62, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x0b, 0x6e, 0x65, 0x61, 0x72, 403 | 0x64, 0x62, 0x76, 0x32, 0x2e, 0x70, 0x62, 0x22, 0x48, 0x0a, 0x0a, 0x41, 0x64, 0x64, 0x52, 0x65, 404 | 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 405 | 0x04, 0x52, 0x02, 0x69, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x74, 0x61, 0x67, 0x6c, 0x69, 0x73, 0x74, 406 | 0x18, 0x02, 0x20, 0x03, 0x28, 0x09, 0x52, 0x07, 0x74, 0x61, 0x67, 0x6c, 0x69, 0x73, 0x74, 0x12, 407 | 0x10, 0x0a, 0x03, 0x70, 0x6f, 0x70, 0x18, 0x03, 0x20, 0x01, 0x28, 0x04, 0x52, 0x03, 0x70, 0x6f, 408 | 0x70, 0x22, 0x1f, 0x0a, 0x0d, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 409 | 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x04, 0x52, 0x02, 410 | 0x69, 0x64, 0x22, 0x0e, 0x0a, 0x0c, 0x4e, 0x6f, 0x6e, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 411 | 0x73, 0x65, 0x22, 0x4c, 0x0a, 0x0c, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x71, 0x75, 0x65, 412 | 0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x74, 0x61, 0x67, 0x6c, 0x69, 0x73, 0x74, 0x18, 0x05, 0x20, 413 | 0x03, 0x28, 0x09, 0x52, 0x07, 0x74, 0x61, 0x67, 0x6c, 0x69, 0x73, 0x74, 0x12, 0x0c, 0x0a, 0x01, 414 | 0x6b, 0x18, 0x06, 0x20, 0x01, 0x28, 0x03, 0x52, 0x01, 0x6b, 0x12, 0x14, 0x0a, 0x05, 0x64, 0x72, 415 | 0x69, 0x66, 0x74, 0x18, 0x07, 0x20, 0x01, 0x28, 0x01, 0x52, 0x05, 0x64, 0x72, 0x69, 0x66, 0x74, 416 | 0x22, 0x32, 0x0a, 0x04, 0x49, 0x74, 0x65, 0x6d, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x08, 417 | 0x20, 0x01, 0x28, 0x04, 0x52, 0x02, 0x69, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x64, 0x69, 0x73, 0x74, 418 | 0x61, 0x6e, 0x63, 0x65, 0x18, 0x09, 0x20, 0x01, 0x28, 0x02, 0x52, 0x08, 0x64, 0x69, 0x73, 0x74, 419 | 0x61, 0x6e, 0x63, 0x65, 0x22, 0x38, 0x0a, 0x0d, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 420 | 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x27, 0x0a, 0x05, 0x69, 0x74, 0x65, 0x6d, 0x73, 0x18, 0x0a, 421 | 0x20, 0x03, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 0x62, 0x76, 0x32, 0x2e, 422 | 0x70, 0x62, 0x2e, 0x49, 0x74, 0x65, 0x6d, 0x52, 0x05, 0x69, 0x74, 0x65, 0x6d, 0x73, 0x22, 0x46, 423 | 0x0a, 0x10, 0x51, 0x75, 0x65, 0x72, 0x79, 0x42, 0x79, 0x49, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 424 | 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x0b, 0x20, 0x01, 0x28, 0x04, 0x52, 0x02, 425 | 0x69, 0x64, 0x12, 0x0c, 0x0a, 0x01, 0x6b, 0x18, 0x0c, 0x20, 0x01, 0x28, 0x03, 0x52, 0x01, 0x6b, 426 | 0x12, 0x14, 0x0a, 0x05, 0x64, 0x72, 0x69, 0x66, 0x74, 0x18, 0x0d, 0x20, 0x01, 0x28, 0x01, 0x52, 427 | 0x05, 0x64, 0x72, 0x69, 0x66, 0x74, 0x32, 0x9b, 0x02, 0x0a, 0x0d, 0x4e, 0x65, 0x61, 0x72, 0x44, 428 | 0x42, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x3b, 0x0a, 0x03, 0x41, 0x64, 0x64, 0x12, 429 | 0x17, 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 0x62, 0x76, 0x32, 0x2e, 0x70, 0x62, 0x2e, 0x41, 0x64, 430 | 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x19, 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 431 | 0x62, 0x76, 0x32, 0x2e, 0x70, 0x62, 0x2e, 0x4e, 0x6f, 0x6e, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 432 | 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x40, 0x0a, 0x05, 0x51, 0x75, 0x65, 0x72, 0x79, 0x12, 0x19, 433 | 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 0x62, 0x76, 0x32, 0x2e, 0x70, 0x62, 0x2e, 0x51, 0x75, 0x65, 434 | 0x72, 0x79, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x6e, 0x65, 0x61, 0x72, 435 | 0x64, 0x62, 0x76, 0x32, 0x2e, 0x70, 0x62, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 436 | 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x48, 0x0a, 0x09, 0x51, 0x75, 0x65, 0x72, 0x79, 437 | 0x42, 0x79, 0x49, 0x64, 0x12, 0x1d, 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 0x62, 0x76, 0x32, 0x2e, 438 | 0x70, 0x62, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x42, 0x79, 0x49, 0x64, 0x52, 0x65, 0x71, 0x75, 439 | 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 0x62, 0x76, 0x32, 0x2e, 0x70, 440 | 0x62, 0x2e, 0x51, 0x75, 0x65, 0x72, 0x79, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 441 | 0x00, 0x12, 0x41, 0x0a, 0x06, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x12, 0x1a, 0x2e, 0x6e, 0x65, 442 | 0x61, 0x72, 0x64, 0x62, 0x76, 0x32, 0x2e, 0x70, 0x62, 0x2e, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 443 | 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x19, 0x2e, 0x6e, 0x65, 0x61, 0x72, 0x64, 0x62, 444 | 0x76, 0x32, 0x2e, 0x70, 0x62, 0x2e, 0x4e, 0x6f, 0x6e, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 445 | 0x73, 0x65, 0x22, 0x00, 0x42, 0x06, 0x5a, 0x04, 0x2e, 0x3b, 0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 446 | 0x6f, 0x74, 0x6f, 0x33, 447 | } 448 | 449 | var ( 450 | file_pb_proto_rawDescOnce sync.Once 451 | file_pb_proto_rawDescData = file_pb_proto_rawDesc 452 | ) 453 | 454 | func file_pb_proto_rawDescGZIP() []byte { 455 | file_pb_proto_rawDescOnce.Do(func() { 456 | file_pb_proto_rawDescData = protoimpl.X.CompressGZIP(file_pb_proto_rawDescData) 457 | }) 458 | return file_pb_proto_rawDescData 459 | } 460 | 461 | var file_pb_proto_msgTypes = make([]protoimpl.MessageInfo, 7) 462 | var file_pb_proto_goTypes = []interface{}{ 463 | (*AddRequest)(nil), // 0: neardbv2.pb.AddRequest 464 | (*RemoveRequest)(nil), // 1: neardbv2.pb.RemoveRequest 465 | (*NoneResponse)(nil), // 2: neardbv2.pb.NoneResponse 466 | (*QueryRequest)(nil), // 3: neardbv2.pb.QueryRequest 467 | (*Item)(nil), // 4: neardbv2.pb.Item 468 | (*QueryResponse)(nil), // 5: neardbv2.pb.QueryResponse 469 | (*QueryByIdRequest)(nil), // 6: neardbv2.pb.QueryByIdRequest 470 | } 471 | var file_pb_proto_depIdxs = []int32{ 472 | 4, // 0: neardbv2.pb.QueryResponse.items:type_name -> neardbv2.pb.Item 473 | 0, // 1: neardbv2.pb.NearDBService.Add:input_type -> neardbv2.pb.AddRequest 474 | 3, // 2: neardbv2.pb.NearDBService.Query:input_type -> neardbv2.pb.QueryRequest 475 | 6, // 3: neardbv2.pb.NearDBService.QueryById:input_type -> neardbv2.pb.QueryByIdRequest 476 | 1, // 4: neardbv2.pb.NearDBService.Remove:input_type -> neardbv2.pb.RemoveRequest 477 | 2, // 5: neardbv2.pb.NearDBService.Add:output_type -> neardbv2.pb.NoneResponse 478 | 5, // 6: neardbv2.pb.NearDBService.Query:output_type -> neardbv2.pb.QueryResponse 479 | 5, // 7: neardbv2.pb.NearDBService.QueryById:output_type -> neardbv2.pb.QueryResponse 480 | 2, // 8: neardbv2.pb.NearDBService.Remove:output_type -> neardbv2.pb.NoneResponse 481 | 5, // [5:9] is the sub-list for method output_type 482 | 1, // [1:5] is the sub-list for method input_type 483 | 1, // [1:1] is the sub-list for extension type_name 484 | 1, // [1:1] is the sub-list for extension extendee 485 | 0, // [0:1] is the sub-list for field type_name 486 | } 487 | 488 | func init() { file_pb_proto_init() } 489 | func file_pb_proto_init() { 490 | if File_pb_proto != nil { 491 | return 492 | } 493 | if !protoimpl.UnsafeEnabled { 494 | file_pb_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} { 495 | switch v := v.(*AddRequest); i { 496 | case 0: 497 | return &v.state 498 | case 1: 499 | return &v.sizeCache 500 | case 2: 501 | return &v.unknownFields 502 | default: 503 | return nil 504 | } 505 | } 506 | file_pb_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { 507 | switch v := v.(*RemoveRequest); i { 508 | case 0: 509 | return &v.state 510 | case 1: 511 | return &v.sizeCache 512 | case 2: 513 | return &v.unknownFields 514 | default: 515 | return nil 516 | } 517 | } 518 | file_pb_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} { 519 | switch v := v.(*NoneResponse); i { 520 | case 0: 521 | return &v.state 522 | case 1: 523 | return &v.sizeCache 524 | case 2: 525 | return &v.unknownFields 526 | default: 527 | return nil 528 | } 529 | } 530 | file_pb_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} { 531 | switch v := v.(*QueryRequest); i { 532 | case 0: 533 | return &v.state 534 | case 1: 535 | return &v.sizeCache 536 | case 2: 537 | return &v.unknownFields 538 | default: 539 | return nil 540 | } 541 | } 542 | file_pb_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} { 543 | switch v := v.(*Item); i { 544 | case 0: 545 | return &v.state 546 | case 1: 547 | return &v.sizeCache 548 | case 2: 549 | return &v.unknownFields 550 | default: 551 | return nil 552 | } 553 | } 554 | file_pb_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} { 555 | switch v := v.(*QueryResponse); i { 556 | case 0: 557 | return &v.state 558 | case 1: 559 | return &v.sizeCache 560 | case 2: 561 | return &v.unknownFields 562 | default: 563 | return nil 564 | } 565 | } 566 | file_pb_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} { 567 | switch v := v.(*QueryByIdRequest); i { 568 | case 0: 569 | return &v.state 570 | case 1: 571 | return &v.sizeCache 572 | case 2: 573 | return &v.unknownFields 574 | default: 575 | return nil 576 | } 577 | } 578 | } 579 | type x struct{} 580 | out := protoimpl.TypeBuilder{ 581 | File: protoimpl.DescBuilder{ 582 | GoPackagePath: reflect.TypeOf(x{}).PkgPath(), 583 | RawDescriptor: file_pb_proto_rawDesc, 584 | NumEnums: 0, 585 | NumMessages: 7, 586 | NumExtensions: 0, 587 | NumServices: 1, 588 | }, 589 | GoTypes: file_pb_proto_goTypes, 590 | DependencyIndexes: file_pb_proto_depIdxs, 591 | MessageInfos: file_pb_proto_msgTypes, 592 | }.Build() 593 | File_pb_proto = out.File 594 | file_pb_proto_rawDesc = nil 595 | file_pb_proto_goTypes = nil 596 | file_pb_proto_depIdxs = nil 597 | } 598 | -------------------------------------------------------------------------------- /common/database/drivers/pb/pb.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | option go_package = ".;pb"; 3 | package neardbv2.pb; 4 | 5 | service NearDBService { 6 | rpc Add (AddRequest) returns (NoneResponse) {} 7 | rpc Query (QueryRequest) returns (QueryResponse) {} 8 | rpc QueryById (QueryByIdRequest) returns (QueryResponse) {} 9 | rpc Remove (RemoveRequest) returns (NoneResponse) {} 10 | } 11 | 12 | message AddRequest { 13 | uint64 id = 1; 14 | repeated string taglist = 2; 15 | uint64 pop = 3; 16 | } 17 | 18 | message RemoveRequest { 19 | uint64 id = 4; 20 | } 21 | 22 | message NoneResponse { 23 | } 24 | 25 | message QueryRequest { 26 | repeated string taglist = 5; 27 | int64 k = 6; 28 | double drift = 7; 29 | } 30 | 31 | message Item { 32 | uint64 id = 8; 33 | float distance = 9; 34 | } 35 | 36 | message QueryResponse { 37 | repeated Item items = 10; 38 | } 39 | 40 | message QueryByIdRequest { 41 | uint64 id = 11; 42 | int64 k = 12; 43 | double drift = 13; 44 | } -------------------------------------------------------------------------------- /common/database/drivers/pb/pb_grpc.pb.go: -------------------------------------------------------------------------------- 1 | // Code generated by protoc-gen-go-grpc. DO NOT EDIT. 2 | // versions: 3 | // - protoc-gen-go-grpc v1.2.0 4 | // - protoc v3.21.2 5 | // source: pb.proto 6 | 7 | package pb 8 | 9 | import ( 10 | context "context" 11 | grpc "google.golang.org/grpc" 12 | codes "google.golang.org/grpc/codes" 13 | status "google.golang.org/grpc/status" 14 | ) 15 | 16 | // This is a compile-time assertion to ensure that this generated file 17 | // is compatible with the grpc package it is being compiled against. 18 | // Requires gRPC-Go v1.32.0 or later. 19 | const _ = grpc.SupportPackageIsVersion7 20 | 21 | // NearDBServiceClient is the client API for NearDBService service. 22 | // 23 | // For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. 24 | type NearDBServiceClient interface { 25 | Add(ctx context.Context, in *AddRequest, opts ...grpc.CallOption) (*NoneResponse, error) 26 | Query(ctx context.Context, in *QueryRequest, opts ...grpc.CallOption) (*QueryResponse, error) 27 | QueryById(ctx context.Context, in *QueryByIdRequest, opts ...grpc.CallOption) (*QueryResponse, error) 28 | Remove(ctx context.Context, in *RemoveRequest, opts ...grpc.CallOption) (*NoneResponse, error) 29 | } 30 | 31 | type nearDBServiceClient struct { 32 | cc grpc.ClientConnInterface 33 | } 34 | 35 | func NewNearDBServiceClient(cc grpc.ClientConnInterface) NearDBServiceClient { 36 | return &nearDBServiceClient{cc} 37 | } 38 | 39 | func (c *nearDBServiceClient) Add(ctx context.Context, in *AddRequest, opts ...grpc.CallOption) (*NoneResponse, error) { 40 | out := new(NoneResponse) 41 | err := c.cc.Invoke(ctx, "/neardbv2.pb.NearDBService/Add", in, out, opts...) 42 | if err != nil { 43 | return nil, err 44 | } 45 | return out, nil 46 | } 47 | 48 | func (c *nearDBServiceClient) Query(ctx context.Context, in *QueryRequest, opts ...grpc.CallOption) (*QueryResponse, error) { 49 | out := new(QueryResponse) 50 | err := c.cc.Invoke(ctx, "/neardbv2.pb.NearDBService/Query", in, out, opts...) 51 | if err != nil { 52 | return nil, err 53 | } 54 | return out, nil 55 | } 56 | 57 | func (c *nearDBServiceClient) QueryById(ctx context.Context, in *QueryByIdRequest, opts ...grpc.CallOption) (*QueryResponse, error) { 58 | out := new(QueryResponse) 59 | err := c.cc.Invoke(ctx, "/neardbv2.pb.NearDBService/QueryById", in, out, opts...) 60 | if err != nil { 61 | return nil, err 62 | } 63 | return out, nil 64 | } 65 | 66 | func (c *nearDBServiceClient) Remove(ctx context.Context, in *RemoveRequest, opts ...grpc.CallOption) (*NoneResponse, error) { 67 | out := new(NoneResponse) 68 | err := c.cc.Invoke(ctx, "/neardbv2.pb.NearDBService/Remove", in, out, opts...) 69 | if err != nil { 70 | return nil, err 71 | } 72 | return out, nil 73 | } 74 | 75 | // NearDBServiceServer is the server API for NearDBService service. 76 | // All implementations must embed UnimplementedNearDBServiceServer 77 | // for forward compatibility 78 | type NearDBServiceServer interface { 79 | Add(context.Context, *AddRequest) (*NoneResponse, error) 80 | Query(context.Context, *QueryRequest) (*QueryResponse, error) 81 | QueryById(context.Context, *QueryByIdRequest) (*QueryResponse, error) 82 | Remove(context.Context, *RemoveRequest) (*NoneResponse, error) 83 | mustEmbedUnimplementedNearDBServiceServer() 84 | } 85 | 86 | // UnimplementedNearDBServiceServer must be embedded to have forward compatible implementations. 87 | type UnimplementedNearDBServiceServer struct { 88 | } 89 | 90 | func (UnimplementedNearDBServiceServer) Add(context.Context, *AddRequest) (*NoneResponse, error) { 91 | return nil, status.Errorf(codes.Unimplemented, "method Add not implemented") 92 | } 93 | func (UnimplementedNearDBServiceServer) Query(context.Context, *QueryRequest) (*QueryResponse, error) { 94 | return nil, status.Errorf(codes.Unimplemented, "method Query not implemented") 95 | } 96 | func (UnimplementedNearDBServiceServer) QueryById(context.Context, *QueryByIdRequest) (*QueryResponse, error) { 97 | return nil, status.Errorf(codes.Unimplemented, "method QueryById not implemented") 98 | } 99 | func (UnimplementedNearDBServiceServer) Remove(context.Context, *RemoveRequest) (*NoneResponse, error) { 100 | return nil, status.Errorf(codes.Unimplemented, "method Remove not implemented") 101 | } 102 | func (UnimplementedNearDBServiceServer) mustEmbedUnimplementedNearDBServiceServer() {} 103 | 104 | // UnsafeNearDBServiceServer may be embedded to opt out of forward compatibility for this service. 105 | // Use of this interface is not recommended, as added methods to NearDBServiceServer will 106 | // result in compilation errors. 107 | type UnsafeNearDBServiceServer interface { 108 | mustEmbedUnimplementedNearDBServiceServer() 109 | } 110 | 111 | func RegisterNearDBServiceServer(s grpc.ServiceRegistrar, srv NearDBServiceServer) { 112 | s.RegisterService(&NearDBService_ServiceDesc, srv) 113 | } 114 | 115 | func _NearDBService_Add_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { 116 | in := new(AddRequest) 117 | if err := dec(in); err != nil { 118 | return nil, err 119 | } 120 | if interceptor == nil { 121 | return srv.(NearDBServiceServer).Add(ctx, in) 122 | } 123 | info := &grpc.UnaryServerInfo{ 124 | Server: srv, 125 | FullMethod: "/neardbv2.pb.NearDBService/Add", 126 | } 127 | handler := func(ctx context.Context, req interface{}) (interface{}, error) { 128 | return srv.(NearDBServiceServer).Add(ctx, req.(*AddRequest)) 129 | } 130 | return interceptor(ctx, in, info, handler) 131 | } 132 | 133 | func _NearDBService_Query_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { 134 | in := new(QueryRequest) 135 | if err := dec(in); err != nil { 136 | return nil, err 137 | } 138 | if interceptor == nil { 139 | return srv.(NearDBServiceServer).Query(ctx, in) 140 | } 141 | info := &grpc.UnaryServerInfo{ 142 | Server: srv, 143 | FullMethod: "/neardbv2.pb.NearDBService/Query", 144 | } 145 | handler := func(ctx context.Context, req interface{}) (interface{}, error) { 146 | return srv.(NearDBServiceServer).Query(ctx, req.(*QueryRequest)) 147 | } 148 | return interceptor(ctx, in, info, handler) 149 | } 150 | 151 | func _NearDBService_QueryById_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { 152 | in := new(QueryByIdRequest) 153 | if err := dec(in); err != nil { 154 | return nil, err 155 | } 156 | if interceptor == nil { 157 | return srv.(NearDBServiceServer).QueryById(ctx, in) 158 | } 159 | info := &grpc.UnaryServerInfo{ 160 | Server: srv, 161 | FullMethod: "/neardbv2.pb.NearDBService/QueryById", 162 | } 163 | handler := func(ctx context.Context, req interface{}) (interface{}, error) { 164 | return srv.(NearDBServiceServer).QueryById(ctx, req.(*QueryByIdRequest)) 165 | } 166 | return interceptor(ctx, in, info, handler) 167 | } 168 | 169 | func _NearDBService_Remove_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { 170 | in := new(RemoveRequest) 171 | if err := dec(in); err != nil { 172 | return nil, err 173 | } 174 | if interceptor == nil { 175 | return srv.(NearDBServiceServer).Remove(ctx, in) 176 | } 177 | info := &grpc.UnaryServerInfo{ 178 | Server: srv, 179 | FullMethod: "/neardbv2.pb.NearDBService/Remove", 180 | } 181 | handler := func(ctx context.Context, req interface{}) (interface{}, error) { 182 | return srv.(NearDBServiceServer).Remove(ctx, req.(*RemoveRequest)) 183 | } 184 | return interceptor(ctx, in, info, handler) 185 | } 186 | 187 | // NearDBService_ServiceDesc is the grpc.ServiceDesc for NearDBService service. 188 | // It's only intended for direct use with grpc.RegisterService, 189 | // and not to be introspected or modified (even as a copy) 190 | var NearDBService_ServiceDesc = grpc.ServiceDesc{ 191 | ServiceName: "neardbv2.pb.NearDBService", 192 | HandlerType: (*NearDBServiceServer)(nil), 193 | Methods: []grpc.MethodDesc{ 194 | { 195 | MethodName: "Add", 196 | Handler: _NearDBService_Add_Handler, 197 | }, 198 | { 199 | MethodName: "Query", 200 | Handler: _NearDBService_Query_Handler, 201 | }, 202 | { 203 | MethodName: "QueryById", 204 | Handler: _NearDBService_QueryById_Handler, 205 | }, 206 | { 207 | MethodName: "Remove", 208 | Handler: _NearDBService_Remove_Handler, 209 | }, 210 | }, 211 | Streams: []grpc.StreamDesc{}, 212 | Metadata: "pb.proto", 213 | } 214 | -------------------------------------------------------------------------------- /common/database/drivers/rabbitmq.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import ( 4 | "log" 5 | "sync" 6 | 7 | "github.com/ShugetsuSoft/pixivel-back/common/models" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 9 | amqp "github.com/rabbitmq/amqp091-go" 10 | ) 11 | 12 | type RabbitMQ struct { 13 | conn *amqp.Connection 14 | chann *amqp.Channel 15 | closesign chan error 16 | reconnectsignal *sync.Cond 17 | } 18 | 19 | func NewRabbitMQ(uri string) (*RabbitMQ, error) { 20 | rbmq := &RabbitMQ{} 21 | reconnectsignal := sync.NewCond(&sync.Mutex{}) 22 | rbmq.reconnectsignal = reconnectsignal 23 | 24 | connect := func() (chan *amqp.Error, error) { 25 | conn, err := amqp.Dial(uri) 26 | if err != nil { 27 | return nil, err 28 | } 29 | rbmq.conn = conn 30 | ch, err := conn.Channel() 31 | ch.Qos(100, 0, false) 32 | if err != nil { 33 | return nil, err 34 | } 35 | rbmq.chann = ch 36 | reconnectsignal.Broadcast() 37 | closesignal := make(chan *amqp.Error, 1) 38 | return conn.NotifyClose(closesignal), nil 39 | } 40 | 41 | closechan, err := connect() 42 | if err != nil { 43 | return nil, err 44 | } 45 | rbmq.closesign = make(chan error) 46 | 47 | go func() { 48 | for { 49 | select { 50 | case e := <-closechan: 51 | telemetry.Log(telemetry.Label{"pos": "rabbitmq"}, e.Error()) 52 | closechan, err = connect() 53 | if err != nil { 54 | log.Fatal(err) 55 | } 56 | case <-rbmq.closesign: 57 | rbmq.conn.Close() 58 | rbmq.chann.Close() 59 | close(rbmq.closesign) 60 | close(closechan) 61 | return 62 | } 63 | } 64 | }() 65 | 66 | return rbmq, nil 67 | } 68 | 69 | func (mq *RabbitMQ) SetQos(prefetchCount int, prefetchSize int, global bool) error { 70 | return mq.chann.Qos(prefetchCount, prefetchSize, global) 71 | } 72 | 73 | func (mq *RabbitMQ) QueueDeclare(name string) (amqp.Queue, error) { 74 | queue, err := mq.chann.QueueDeclare( 75 | name, 76 | true, // durable 77 | false, // delete when unused 78 | false, // exclusive 79 | false, // no-wait 80 | nil, // arguments 81 | ) 82 | return queue, err 83 | } 84 | 85 | func (mq *RabbitMQ) ExchangeDeclare(name string, kind string) error { 86 | return mq.chann.ExchangeDeclare( 87 | name, 88 | kind, 89 | true, false, false, false, nil) 90 | } 91 | 92 | func (mq *RabbitMQ) QueueBindExchange(name, key, exchange string) { 93 | mq.chann.QueueBind(name, key, exchange, false, nil) 94 | } 95 | 96 | func (mq *RabbitMQ) Publish(name string, message models.MQMessage) error { 97 | return mq.chann.Publish( 98 | "", // exchange 99 | name, // routing key 100 | false, // mandatory 101 | false, 102 | amqp.Publishing{ 103 | DeliveryMode: amqp.Persistent, 104 | Body: message.Data, 105 | Priority: message.Priority, 106 | }, 107 | ) 108 | } 109 | 110 | func (mq *RabbitMQ) PublishToExchange(name string, message models.MQMessage) error { 111 | return mq.chann.Publish( 112 | name, // exchange 113 | name, // routing key 114 | false, // mandatory 115 | false, 116 | amqp.Publishing{ 117 | DeliveryMode: amqp.Persistent, 118 | Body: message.Data, 119 | Priority: message.Priority, 120 | }, 121 | ) 122 | } 123 | 124 | func (mq *RabbitMQ) Consume(name string) (<-chan models.MQMessage, error) { 125 | reschan := make(chan models.MQMessage, 50) 126 | mq.reconnectsignal.L.Lock() 127 | if mq.chann.IsClosed() { 128 | mq.reconnectsignal.Wait() 129 | } 130 | mq.reconnectsignal.L.Unlock() 131 | 132 | oniichan, err := mq.chann.Consume( 133 | name, 134 | "", 135 | false, 136 | false, 137 | false, 138 | false, 139 | nil, 140 | ) 141 | go func() { 142 | defer close(reschan) 143 | for onii := range oniichan { 144 | reschan <- models.MQMessage{ 145 | Data: onii.Body, 146 | Tag: onii.DeliveryTag, 147 | Priority: onii.Priority, 148 | } 149 | } 150 | }() 151 | return reschan, err 152 | } 153 | 154 | func (mq *RabbitMQ) Ack(tag uint64) error { 155 | return mq.chann.Ack(tag, false) 156 | } 157 | 158 | func (mq *RabbitMQ) Reject(tag uint64) error { 159 | return mq.chann.Reject(tag, false) 160 | } 161 | 162 | func (mq *RabbitMQ) Close() { 163 | mq.closesign <- nil 164 | } 165 | -------------------------------------------------------------------------------- /common/database/drivers/redis.go: -------------------------------------------------------------------------------- 1 | package drivers 2 | 3 | import ( 4 | "context" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 6 | "time" 7 | 8 | "github.com/gomodule/redigo/redis" 9 | ) 10 | 11 | type RedisPool struct { 12 | RedisPool *redis.Pool 13 | } 14 | 15 | func NewRedisPool(redisURL string) *RedisPool { 16 | return &RedisPool{ 17 | RedisPool: &redis.Pool{ 18 | MaxIdle: 100, 19 | MaxActive: 0, 20 | Wait: true, 21 | IdleTimeout: time.Duration(100) * time.Second, 22 | Dial: func() (redis.Conn, error) { 23 | c, err := redis.DialURL(redisURL) 24 | if err != nil { 25 | return nil, err 26 | } 27 | return c, err 28 | }, 29 | }, 30 | } 31 | } 32 | 33 | type RedisClient struct { 34 | conn redis.Conn 35 | } 36 | 37 | func (red *RedisPool) NewRedisClient() *RedisClient { 38 | conn := red.RedisPool.Get() 39 | return &RedisClient{ 40 | conn: conn, 41 | } 42 | } 43 | 44 | func NewRedisClient(redisURL string) (*RedisClient, error) { 45 | conn, err := redis.DialURL(redisURL) 46 | if err != nil { 47 | return nil, err 48 | } 49 | return &RedisClient{ 50 | conn: conn, 51 | }, nil 52 | } 53 | 54 | func (red *RedisClient) IsExist(err error) bool { 55 | return err == nil 56 | } 57 | 58 | func (red *RedisClient) IsError(err error) bool { 59 | return err != redis.ErrNil 60 | } 61 | 62 | func (red *RedisClient) GetValue(key string) (string, error) { 63 | value, err := redis.String(red.conn.Do("GET", key)) 64 | return value, err 65 | } 66 | 67 | func (red *RedisClient) GetValueBytes(key string) ([]byte, error) { 68 | value, err := redis.Bytes(red.conn.Do("GET", key)) 69 | return value, err 70 | } 71 | 72 | func (red *RedisClient) SetValue(key string, value string) error { 73 | _, err := red.conn.Do("SET", key, value) 74 | return err 75 | } 76 | 77 | func (red *RedisClient) SetValueBytes(key string, value []byte) error { 78 | _, err := red.conn.Do("SET", key, value) 79 | return err 80 | } 81 | 82 | func (red *RedisClient) DeleteValue(key string) error { 83 | _, err := red.conn.Do("DEL", key) 84 | return err 85 | } 86 | 87 | func (red *RedisClient) SetValueExpire(key string, value string, exp int) error { 88 | _, err := red.conn.Do("SETEX", key, exp, value) 89 | return err 90 | } 91 | 92 | func (red *RedisClient) SetValueBytesExpire(key string, value []byte, exp int) error { 93 | _, err := red.conn.Do("SETEX", key, exp, value) 94 | return err 95 | } 96 | 97 | func (red *RedisClient) ClearDB() error { 98 | _, err := red.conn.Do("FLUSHDB") 99 | return err 100 | } 101 | 102 | func (red *RedisClient) SetExpire(key string, exp string) error { 103 | _, err := red.conn.Do("EXPIRE", key, exp) 104 | return err 105 | } 106 | func (red *RedisClient) GetExpire(key string) (string, error) { 107 | value, err := redis.String(red.conn.Do("TTL", key)) 108 | if err != nil { 109 | if err == redis.ErrNil { 110 | return "nil", err 111 | } 112 | return "", err 113 | } 114 | return value, nil 115 | } 116 | 117 | func (red *RedisClient) KeyExist(key string) (bool, error) { 118 | exist, err := redis.Bool(red.conn.Do("EXISTS", key)) 119 | if err != nil { 120 | return false, err 121 | } else { 122 | return exist, nil 123 | } 124 | } 125 | 126 | func (red *RedisClient) Increase(key string) (int, error) { 127 | return redis.Int(red.conn.Do("INCR", key)) 128 | } 129 | 130 | func (red *RedisClient) IncreaseBy(key string, num uint) (int, error) { 131 | return redis.Int(red.conn.Do("INCRBY", key, num)) 132 | } 133 | 134 | func (red *RedisClient) Decrease(key string) (int, error) { 135 | return redis.Int(red.conn.Do("DECR", key)) 136 | } 137 | 138 | func (red *RedisClient) DecreaseBy(key string, num uint) (int, error) { 139 | return redis.Int(red.conn.Do("DECRBY", key, num)) 140 | } 141 | 142 | func (red *RedisClient) Publish(key string, message string) error { 143 | _, err := red.conn.Do("PUBLISH", key, message) 144 | return err 145 | } 146 | 147 | func (red *RedisClient) Subscribe(ctx context.Context, key string) (chan string, error) { 148 | psc := redis.PubSubConn{Conn: red.conn} 149 | if err := psc.PSubscribe(key); err != nil { 150 | return nil, err 151 | } 152 | msg := make(chan string, 5) 153 | go func() { 154 | for { 155 | select { 156 | case <-ctx.Done(): 157 | close(msg) 158 | return 159 | default: 160 | switch v := psc.Receive().(type) { 161 | case redis.Message: 162 | msg <- utils.StringOut(v.Data) 163 | } 164 | } 165 | } 166 | }() 167 | return msg, nil 168 | } 169 | 170 | func (red *RedisClient) Close() { 171 | red.conn.Close() 172 | } 173 | -------------------------------------------------------------------------------- /common/database/operations/illust.go: -------------------------------------------------------------------------------- 1 | package operations 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "time" 7 | 8 | "github.com/ShugetsuSoft/pixivel-back/common/convert" 9 | "github.com/ShugetsuSoft/pixivel-back/common/models" 10 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 11 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 12 | "github.com/olivere/elastic/v7" 13 | "go.mongodb.org/mongo-driver/bson" 14 | "go.mongodb.org/mongo-driver/mongo" 15 | "go.mongodb.org/mongo-driver/mongo/options" 16 | ) 17 | 18 | func (ops *DatabaseOperations) InsertIllust(ctx context.Context, illust *models.Illust) error { 19 | var err error 20 | is, err := ops.Flt.Exists(config.IllustTableName, utils.Itoa(illust.ID)) 21 | if err != nil { 22 | return err 23 | } 24 | illust.UpdateTime = time.Now() 25 | 26 | if is { 27 | goto REPLACE 28 | } else { 29 | _, err = ops.Cols.Illust.InsertOne(ctx, illust) 30 | 31 | if mongo.IsDuplicateKeyError(err) { 32 | _, err = ops.Flt.Add(config.IllustTableName, utils.Itoa(illust.ID)) 33 | if err != nil { 34 | return err 35 | } 36 | goto REPLACE 37 | } 38 | 39 | if err != nil { 40 | return err 41 | } 42 | 43 | _, err = ops.Flt.Add(config.IllustTableName, utils.Itoa(illust.ID)) 44 | if err != nil { 45 | return err 46 | } 47 | } 48 | 49 | return nil 50 | 51 | REPLACE: 52 | result, err := ops.Cols.Illust.ReplaceOne(ctx, bson.M{"_id": illust.ID}, illust) 53 | if err != nil { 54 | return err 55 | } 56 | if result.MatchedCount == 0 { 57 | _, err = ops.Cols.Illust.InsertOne(ctx, illust) 58 | if err != nil { 59 | return err 60 | } 61 | _, err = ops.Flt.Add(config.IllustTableName, utils.Itoa(illust.ID)) 62 | } 63 | 64 | return err 65 | } 66 | 67 | func (ops *DatabaseOperations) AddIllusts(ctx context.Context, illusts []models.Illust) error { 68 | operations := make([]mongo.WriteModel, len(illusts)) 69 | for i, illust := range illusts { 70 | ins := mongo.NewInsertOneModel() 71 | ins.SetDocument(illust) 72 | operations[i] = ins 73 | } 74 | bulkOption := &options.BulkWriteOptions{} 75 | bulkOption.SetOrdered(false) 76 | bulkOption.SetBypassDocumentValidation(false) 77 | _, err := ops.Cols.Illust.BulkWrite(ctx, operations, bulkOption) 78 | return err 79 | } 80 | 81 | func (ops *DatabaseOperations) InsertIllusts(ctx context.Context, illusts []models.Illust) error { 82 | nowIllusts := make([]interface{}, 0, len(illusts)) 83 | updateoperations := make([]mongo.WriteModel, 0, len(illusts)) 84 | 85 | for _, illust := range illusts { 86 | is, err := ops.Flt.Exists(config.IllustTableName, utils.Itoa(illust.ID)) 87 | if err != nil { 88 | return err 89 | } 90 | if is { 91 | rep := mongo.NewReplaceOneModel() 92 | rep.SetFilter(bson.M{"_id": illust.ID}) 93 | rep.SetReplacement(illust) 94 | rep.SetUpsert(true) 95 | updateoperations = append(updateoperations, rep) 96 | } else { 97 | _, err := ops.Flt.Add(config.IllustTableName, utils.Itoa(illust.ID)) 98 | if err != nil { 99 | return err 100 | } 101 | nowIllusts = append(nowIllusts, illust) 102 | } 103 | } 104 | 105 | if len(nowIllusts) > 0 { 106 | _, err := ops.Cols.Illust.InsertMany(ctx, nowIllusts) 107 | if err != nil { 108 | return err 109 | } 110 | } 111 | 112 | if len(updateoperations) > 0 { 113 | bulkOption := &options.BulkWriteOptions{} 114 | bulkOption.SetOrdered(false) 115 | bulkOption.SetBypassDocumentValidation(false) 116 | _, err := ops.Cols.Illust.BulkWrite(ctx, updateoperations, bulkOption) 117 | return err 118 | } 119 | return nil 120 | } 121 | 122 | func (ops *DatabaseOperations) InsertIllustSearch(ctx context.Context, illust *models.Illust) error { 123 | illustsearch := convert.Illust2IllustSearch(illust) 124 | err := ops.Sc.es.InsertDocument(ctx, config.IllustSearchIndexName, utils.Itoa(illust.ID), illustsearch) 125 | if err != nil { 126 | return err 127 | } 128 | return nil 129 | } 130 | 131 | func (ops *DatabaseOperations) InsertIllustTagNearDB(ctx context.Context, illust *models.Illust) error { 132 | tagset := make([]string, len(illust.Tags)) 133 | for i, tag := range illust.Tags { 134 | tagset[i] = tag.Name 135 | } 136 | 137 | return ops.Sc.ndb.Add(ctx, illust.ID, tagset, illust.Popularity) 138 | } 139 | 140 | func (ops *DatabaseOperations) RecommendIllustsByIllustId(ctx context.Context, illustId uint64, k int, drif float64, resultbanned bool) ([]models.Illust, error) { 141 | items, err := ops.Sc.ndb.QueryById(ctx, illustId, k, drif) 142 | if err != nil { 143 | return nil, err 144 | } 145 | if len(items) < 1 { 146 | illust, err := ops.QueryIllust(ctx, illustId, resultbanned) 147 | if err != nil { 148 | return nil, err 149 | } 150 | tags := make([]string, len(illust.Tags)) 151 | for i, tag := range illust.Tags { 152 | tags[i] = tag.Name 153 | } 154 | err = ops.Sc.ndb.Add(ctx, illust.ID, tags, illust.Popularity) 155 | if err != nil { 156 | return nil, err 157 | } 158 | items, err = ops.Sc.ndb.QueryById(ctx, illustId, k, drif) 159 | if err != nil { 160 | return nil, err 161 | } 162 | } 163 | queryidlist := make([]uint64, 0, len(items)) 164 | for _, item := range items { 165 | if item.Id == illustId { 166 | continue 167 | } 168 | queryidlist = append(queryidlist, item.Id) 169 | } 170 | 171 | illusts, err := ops.QueryIllusts(ctx, queryidlist, resultbanned) 172 | if err != nil { 173 | return nil, err 174 | } 175 | 176 | illustsmap := make(map[uint64]models.Illust) 177 | for _, illust := range illusts { 178 | illustsmap[illust.ID] = illust 179 | } 180 | 181 | res := make([]models.Illust, 0, len(items)) 182 | for _, illustid := range queryidlist { 183 | if _, exist := illustsmap[illustid]; exist { 184 | res = append(res, illustsmap[illustid]) 185 | } 186 | } 187 | 188 | return res, nil 189 | } 190 | 191 | func (ops *DatabaseOperations) QueryIllust(ctx context.Context, illustId uint64, resultbanned bool) (*models.Illust, error) { 192 | is, err := ops.Flt.Exists(config.IllustTableName, utils.Itoa(illustId)) 193 | 194 | if err != nil { 195 | return nil, err 196 | } 197 | 198 | if is { 199 | result := models.Illust{ 200 | Statistic: models.IllustStatistic{}, 201 | Tags: []models.IllustTag{}, 202 | } 203 | query := bson.M{"_id": illustId} 204 | err := ops.Cols.Illust.FindOne(ctx, query).Decode(&result) 205 | if err != nil { 206 | if err == mongo.ErrNoDocuments { 207 | return nil, nil 208 | } else { 209 | return nil, err 210 | } 211 | } 212 | 213 | if result.Banned && !resultbanned { 214 | return nil, models.ErrorItemBanned 215 | } 216 | 217 | return &result, err 218 | } 219 | return nil, nil 220 | } 221 | 222 | func (ops *DatabaseOperations) QueryIllusts(ctx context.Context, illustIds []uint64, resultbanned bool) ([]models.Illust, error) { 223 | query := bson.M{ 224 | "_id": bson.M{"$in": illustIds}, 225 | "banned": false, 226 | } 227 | if resultbanned { 228 | query["banned"] = true 229 | } 230 | cursor, err := ops.Cols.Illust.Find(ctx, query) 231 | defer cursor.Close(ctx) 232 | if err != nil { 233 | if err == mongo.ErrNoDocuments { 234 | return nil, nil 235 | } else { 236 | return nil, err 237 | } 238 | } 239 | 240 | illusts := make([]models.Illust, 0, len(illustIds)) 241 | 242 | for cursor.Next(ctx) { 243 | result := models.Illust{ 244 | Statistic: models.IllustStatistic{}, 245 | Tags: []models.IllustTag{}, 246 | } 247 | err := cursor.Decode(&result) 248 | if err != nil { 249 | return nil, err 250 | } 251 | 252 | illusts = append(illusts, result) 253 | } 254 | 255 | return illusts, err 256 | } 257 | 258 | func (ops *DatabaseOperations) SearchIllustSuggest(ctx context.Context, keyword string) ([]string, error) { 259 | source := elastic.NewSearchSource(). 260 | Suggester(ops.Sc.es.Suggest("title-completion-suggest").Field("title.suggest").Text(keyword).Fuzziness(2).Analyzer("kuromoji").SkipDuplicates(true)). 261 | Suggester(ops.Sc.es.Suggest("alt_title-completion-suggest").Field("alt_title.suggest").Text(keyword).Fuzziness(3).Analyzer("kuromoji").SkipDuplicates(true)). 262 | Suggester(ops.Sc.es.Suggest("name-completion-suggest").Field("tags.name.suggest").Text(keyword).Fuzziness(2).Analyzer("kuromoji").SkipDuplicates(true)). 263 | Suggester(ops.Sc.es.Suggest("trans-completion-suggest").Field("tags.translation.suggest").Text(keyword).Fuzziness(2).Analyzer("smartcn").SkipDuplicates(true)). 264 | FetchSource(false) 265 | query := ops.Sc.es.Search(config.IllustSearchIndexName).Source(nil). 266 | SearchSource(source) 267 | 268 | results, err := ops.Sc.es.DoSearch(ctx, query) 269 | if err != nil { 270 | return nil, err 271 | } 272 | suggests := append(append(append(results.Suggest["title-completion-suggest"][0].Options, results.Suggest["alt_title-completion-suggest"][0].Options...), results.Suggest["name-completion-suggest"][0].Options...), results.Suggest["trans-completion-suggest"][0].Options...) 273 | res := make([]string, 0, len(suggests)) 274 | uniqmap := make(map[string]bool, len(suggests)) 275 | for _, suggest := range suggests { 276 | if uniqmap[suggest.Text] == false { 277 | uniqmap[suggest.Text] = true 278 | res = append(res, suggest.Text) 279 | } 280 | } 281 | return res, nil 282 | } 283 | 284 | func (ops *DatabaseOperations) SearchTagSuggest(ctx context.Context, keyword string) ([]models.IllustTag, error) { 285 | source := elastic.NewSearchSource(). 286 | Suggester(ops.Sc.es.Suggest("name-completion-suggest").Field("tags.name.suggest").Text(keyword).Fuzziness(2).Analyzer("kuromoji").SkipDuplicates(true)). 287 | Suggester(ops.Sc.es.Suggest("trans-completion-suggest").Field("tags.translation.suggest").Text(keyword).Fuzziness(2).Analyzer("smartcn").SkipDuplicates(true)). 288 | FetchSource(true) 289 | query := ops.Sc.es.Search(config.IllustSearchIndexName).Source(nil). 290 | SearchSource(source) 291 | 292 | results, err := ops.Sc.es.DoSearch(ctx, query) 293 | if err != nil { 294 | return nil, err 295 | } 296 | suggests := append(results.Suggest["name-completion-suggest"][0].Options, results.Suggest["trans-completion-suggest"][0].Options...) 297 | //suggests := results.Suggest["name-completion-suggest"][0].Options 298 | res := make([]models.IllustTag, len(suggests)) 299 | for i, suggest := range suggests { 300 | var tag models.IllustTag 301 | json.Unmarshal(suggest.Source, &tag) 302 | res[i] = tag 303 | } 304 | return res, nil 305 | } 306 | 307 | func (ops *DatabaseOperations) SearchIllust(ctx context.Context, keyword string, page int, limit int, sortpopularity bool, sortdate bool, resultbanned bool) ([]models.Illust, int64, []float64, []*string, error) { 308 | query := ops.Sc.es.Search(config.IllustSearchIndexName). 309 | Query(ops.Sc.es.BoolQuery(). 310 | Should(ops.Sc.es.Query("title", keyword).Analyzer("kuromoji").Boost(4)). 311 | Should(elastic.NewMatchQuery("title.keyword", keyword).Boost(3)). 312 | Should(ops.Sc.es.Query("alt_title", keyword).Analyzer("kuromoji").Boost(2)). 313 | Should(elastic.NewNestedQuery("tags", 314 | ops.Sc.es.BoolQuery().Should(ops.Sc.es.Query("tags.name.fuzzy", keyword).Analyzer("kuromoji").Boost(2)). 315 | Should(ops.Sc.es.Query("tags.translation.fuzzy", keyword).Analyzer("kuromoji").Boost(1))), 316 | ), 317 | ). 318 | Size(limit).From(page * limit). 319 | Highlight(elastic.NewHighlight().Field("title")). 320 | FetchSourceContext(elastic.NewFetchSourceContext(true).Include("_id")).TrackScores(true) 321 | if sortpopularity { 322 | query = query.Sort("popularity", false) 323 | } 324 | if sortdate { 325 | query = query.Sort("create_date", false) 326 | } 327 | query = query.Sort("_score", false) 328 | 329 | results, err := ops.Sc.es.DoSearch(ctx, query) 330 | if err != nil { 331 | return nil, 0, nil, nil, err 332 | } 333 | 334 | if results.Hits.TotalHits.Value > 0 { 335 | scores := make([]float64, 0, len(results.Hits.Hits)) 336 | highlights := make([]*string, 0, len(results.Hits.Hits)) 337 | illustids := make([]uint64, len(results.Hits.Hits)) 338 | for i, hit := range results.Hits.Hits { 339 | illustids[i] = utils.Atoi(hit.Id) 340 | if hit.Score != nil { 341 | scores = append(scores, *hit.Score) 342 | } else { 343 | scores = append(scores, -1) 344 | } 345 | if highl, ok := hit.Highlight["title"]; ok { 346 | highlights = append(highlights, &highl[0]) 347 | } else { 348 | highlights = append(highlights, nil) 349 | } 350 | } 351 | 352 | illusts, err := ops.QueryIllusts(ctx, illustids, resultbanned) 353 | if err != nil { 354 | return nil, 0, nil, nil, err 355 | } 356 | 357 | illustsmap := make(map[uint64]models.Illust) 358 | for _, illust := range illusts { 359 | illustsmap[illust.ID] = illust 360 | } 361 | 362 | result := make([]models.Illust, 0, len(results.Hits.Hits)) 363 | for _, illustid := range illustids { 364 | if _, exist := illustsmap[illustid]; exist { 365 | result = append(result, illustsmap[illustid]) 366 | } 367 | } 368 | 369 | return result, results.Hits.TotalHits.Value, scores, highlights, nil 370 | } else { 371 | return nil, 0, nil, nil, models.ErrorNoResult 372 | } 373 | return nil, 0, nil, nil, err 374 | } 375 | 376 | func (ops *DatabaseOperations) QueryIllustsByTags(ctx context.Context, tags []string, page int, limit int, sortpopularity bool, sortdate bool, resultbanned bool) ([]models.Illust, error) { 377 | /* 378 | { 379 | "query": { 380 | "nested": { 381 | "path": "tags", 382 | "query": { 383 | "bool": { 384 | "must": [ 385 | { 386 | "match": { 387 | "tags.name.fuzzy": "" 388 | } 389 | }, 390 | ] 391 | } 392 | } 393 | } 394 | } 395 | } 396 | */ 397 | tagQuires := make([]elastic.Query, len(tags)) 398 | for i, tag := range tags { 399 | tagQuires[i] = elastic.NewMatchQuery("tags.name.fuzzy", tag) 400 | } 401 | query := ops.Sc.es.Search(config.IllustSearchIndexName).Query( 402 | elastic.NewNestedQuery("tags", elastic.NewBoolQuery().Must(tagQuires...)), 403 | ).Size(limit).From(page * limit).FetchSourceContext(elastic.NewFetchSourceContext(true).Include("_id")) 404 | 405 | if sortpopularity { 406 | query = query.Sort("popularity", false) 407 | } 408 | if sortdate { 409 | query = query.Sort("create_date", false) 410 | } 411 | 412 | results, err := ops.Sc.es.DoSearch(ctx, query) 413 | if err != nil { 414 | return nil, err 415 | } 416 | 417 | if results.Hits.TotalHits.Value > 0 { 418 | illustids := make([]uint64, len(results.Hits.Hits)) 419 | for i, hit := range results.Hits.Hits { 420 | illustids[i] = utils.Atoi(hit.Id) 421 | } 422 | 423 | illusts, err := ops.QueryIllusts(ctx, illustids, resultbanned) 424 | if err != nil { 425 | return nil, err 426 | } 427 | 428 | illustsmap := make(map[uint64]models.Illust) 429 | for _, illust := range illusts { 430 | illustsmap[illust.ID] = illust 431 | } 432 | 433 | result := make([]models.Illust, 0, len(results.Hits.Hits)) 434 | for _, illustid := range illustids { 435 | if _, exist := illustsmap[illustid]; exist { 436 | result = append(result, illustsmap[illustid]) 437 | } 438 | } 439 | 440 | return result, nil 441 | } 442 | return nil, models.ErrorNoResult 443 | } 444 | 445 | func (ops *DatabaseOperations) QueryIllustByUser(ctx context.Context, userId uint64, resultbanned bool) ([]models.Illust, error) { 446 | var results []models.Illust 447 | opts := options.Find().SetSort(bson.M{"createDate": -1}) 448 | query := bson.M{"user": userId, "banned": false} 449 | if resultbanned { 450 | query["banned"] = true 451 | } 452 | cursor, err := ops.Cols.Illust.Find(ctx, query, opts) 453 | if err != nil { 454 | return nil, err 455 | } 456 | 457 | if err = cursor.All(ctx, &results); err != nil { 458 | return nil, err 459 | } 460 | 461 | return results, err 462 | } 463 | 464 | func (ops *DatabaseOperations) QueryIllustByUserWithPage(ctx context.Context, userId uint64, page int64, limit int64, resultbanned bool) ([]models.Illust, error) { 465 | var results []models.Illust 466 | query := bson.M{"user": userId, "banned": false} 467 | if resultbanned { 468 | query["banned"] = true 469 | } 470 | opts := options.Find().SetSort(bson.M{"createDate": -1}).SetLimit(limit).SetSkip(page * limit) 471 | cursor, err := ops.Cols.Illust.Find(ctx, query, opts) 472 | if err != nil { 473 | return nil, err 474 | } 475 | 476 | if err = cursor.All(ctx, &results); err != nil { 477 | return nil, err 478 | } 479 | 480 | return results, err 481 | } 482 | 483 | func (ops *DatabaseOperations) DeleteIllust(ctx context.Context, illustId uint64) error { 484 | _, err := ops.Cols.Illust.DeleteOne(ctx, bson.M{"_id": illustId}) 485 | if err != nil { 486 | if err == mongo.ErrNoDocuments { 487 | return nil 488 | } 489 | return err 490 | } 491 | return err 492 | } 493 | 494 | func (ops *DatabaseOperations) IsIllustExist(illustId uint64) (bool, error) { 495 | is, err := ops.Flt.Exists(config.IllustTableName, utils.Itoa(illustId)) 496 | if err != nil { 497 | return false, err 498 | } 499 | return is, nil 500 | } 501 | -------------------------------------------------------------------------------- /common/database/operations/operations.go: -------------------------------------------------------------------------------- 1 | package operations 2 | 3 | import ( 4 | "context" 5 | "log" 6 | 7 | "go.mongodb.org/mongo-driver/bson" 8 | "go.mongodb.org/mongo-driver/mongo/options" 9 | 10 | "github.com/ShugetsuSoft/pixivel-back/common/database/drivers" 11 | "github.com/ShugetsuSoft/pixivel-back/common/models" 12 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 13 | "go.mongodb.org/mongo-driver/mongo" 14 | ) 15 | 16 | type DatabaseCollections struct { 17 | Illust *mongo.Collection 18 | User *mongo.Collection 19 | Rank *mongo.Collection 20 | Ugoira *mongo.Collection 21 | } 22 | 23 | type DatabaseOperations struct { 24 | Flt models.Filter 25 | Cols *DatabaseCollections 26 | Sc *SearchOperations 27 | } 28 | 29 | type SearchOperations struct { 30 | es *drivers.ElasticSearch 31 | ndb *drivers.NearDB 32 | } 33 | 34 | func NewDatabaseOperations(ctx context.Context, db *drivers.MongoDatabase, filter models.Filter, es *drivers.ElasticSearch, ndb *drivers.NearDB) *DatabaseOperations { 35 | var err error 36 | if es != nil { 37 | err = es.CreateIndex(ctx, config.IllustSearchIndexName, models.IllustSearchMapping) 38 | if err != nil && err != models.ErrorIndexExist { 39 | log.Fatal(err) 40 | } 41 | err = es.CreateIndex(ctx, config.UserSearchIndexName, models.UserSearchMapping) 42 | if err != nil && err != models.ErrorIndexExist { 43 | log.Fatal(err) 44 | } 45 | } 46 | 47 | illustCol := db.Collection(config.IllustTableName) 48 | userCol := db.Collection(config.UserTableName) 49 | rankCol := db.Collection(config.RankTableName) 50 | ugoiraCol := db.Collection(config.UgoiraTableName) 51 | rankCol.Indexes().CreateOne(ctx, mongo.IndexModel{ 52 | Keys: bson.D{{Key: "date", Value: -1}, {Key: "mode", Value: -1}, {Key: "content", Value: -1}}, 53 | Options: options.Index().SetUnique(true), 54 | }) 55 | illustCol.Indexes().CreateOne(ctx, mongo.IndexModel{ 56 | Keys: bson.D{{Key: "tags.name", Value: 1}}, 57 | }) 58 | illustCol.Indexes().CreateOne(ctx, mongo.IndexModel{ 59 | Keys: bson.D{{Key: "user", Value: 1}}, 60 | }) 61 | return &DatabaseOperations{ 62 | Flt: filter, 63 | Cols: &DatabaseCollections{ 64 | Illust: illustCol, 65 | User: userCol, 66 | Rank: rankCol, 67 | Ugoira: ugoiraCol, 68 | }, 69 | Sc: &SearchOperations{ 70 | es: es, 71 | ndb: ndb, 72 | }, 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /common/database/operations/rank.go: -------------------------------------------------------------------------------- 1 | package operations 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/ShugetsuSoft/pixivel-back/common/models" 7 | "go.mongodb.org/mongo-driver/bson" 8 | "go.mongodb.org/mongo-driver/mongo" 9 | ) 10 | 11 | func (ops *DatabaseOperations) InsertRank(ctx context.Context, mode string, date string, content string) (bool, error) { 12 | rank := &models.Rank{ 13 | Date: date, 14 | Mode: mode, 15 | Content: content, 16 | Illusts: []models.RankIllust{}, 17 | } 18 | _, err := ops.Cols.Rank.InsertOne(ctx, rank) 19 | if err != nil { 20 | if mongo.IsDuplicateKeyError(err) { 21 | return false, nil 22 | } 23 | return false, err 24 | } 25 | return true, nil 26 | } 27 | 28 | func (ops *DatabaseOperations) IsRankExist(ctx context.Context, mode string, date string, content string) (bool, error) { 29 | err := ops.Cols.Rank.FindOne(ctx, bson.M{ 30 | "date": date, 31 | "mode": mode, 32 | "content": content, 33 | "$where": "this.illusts.length>0", 34 | }).Err() 35 | if err == mongo.ErrNoDocuments { 36 | return false, nil 37 | } else if err != nil { 38 | return false, err 39 | } 40 | return true, nil 41 | } 42 | 43 | func (ops *DatabaseOperations) AddRankIllusts(ctx context.Context, mode string, date string, content string, illusts []models.RankIllust) error { 44 | filter := bson.M{"date": date, "mode": mode, "content": content} 45 | update := bson.M{ 46 | "$push": bson.M{ 47 | "illusts": bson.M{ 48 | "$each": illusts, 49 | "$sort": bson.M{"rank": 1}, 50 | }, 51 | }, 52 | } 53 | _, err := ops.Cols.Rank.UpdateOne(ctx, filter, update) 54 | return err 55 | } 56 | 57 | func (ops *DatabaseOperations) QueryRankIllusts(ctx context.Context, mode string, date string, content string, page int, limit int) ([]models.RankAggregateResult, error) { 58 | pipeline := mongo.Pipeline{ 59 | {{"$match", bson.D{ 60 | {"date", date}, 61 | {"mode", mode}, 62 | {"content", content}, 63 | }}}, 64 | {{"$unwind", bson.D{ 65 | {"path", "$illusts"}, 66 | }}}, 67 | {{"$skip", page * limit}}, 68 | {{"$limit", limit}}, 69 | {{"$lookup", bson.D{ 70 | {"from", "Illust"}, 71 | {"localField", "illusts.id"}, 72 | {"foreignField", "_id"}, 73 | {"as", "illusts"}, 74 | }}}, 75 | {{"$project", bson.D{ 76 | {"content", bson.D{{"$arrayElemAt", bson.A{"$illusts", 0}}}}, 77 | {"_id", 0}, 78 | }}}, 79 | } 80 | cursor, err := ops.Cols.Rank.Aggregate(ctx, pipeline) 81 | if err != nil { 82 | return nil, err 83 | } 84 | var results []models.RankAggregateResult 85 | if err = cursor.All(ctx, &results); err != nil { 86 | return nil, err 87 | } 88 | return results, nil 89 | } 90 | 91 | func (ops *DatabaseOperations) GetSampleIllusts(ctx context.Context, quality int, limit int, resultbanned bool) ([]models.Illust, error) { 92 | pipeline := mongo.Pipeline{ 93 | {{"$sample", bson.D{ 94 | {"size", limit * limit}, 95 | }}}, 96 | {{"$match", bson.D{ 97 | {"popularity", bson.D{{"$gt", quality}}}, 98 | {"type", 0}, 99 | {"banned", resultbanned}, 100 | }}}, 101 | {{"$sort", bson.D{{ 102 | "popularity", -1, 103 | }}}}, 104 | {{"$limit", limit}}, 105 | } 106 | cursor, err := ops.Cols.Illust.Aggregate(ctx, pipeline) 107 | if err != nil { 108 | return nil, err 109 | } 110 | var results []models.Illust 111 | if err = cursor.All(ctx, &results); err != nil { 112 | return nil, err 113 | } 114 | return results, nil 115 | } 116 | 117 | func (ops *DatabaseOperations) GetSampleUsers(ctx context.Context, limit int, resultbanned bool) ([]models.User, error) { 118 | pipeline := mongo.Pipeline{ 119 | {{"$sample", bson.D{ 120 | {"size", limit}, 121 | }}}, 122 | {{"$match", bson.D{ 123 | {"banned", resultbanned}, 124 | }}}, 125 | } 126 | cursor, err := ops.Cols.User.Aggregate(ctx, pipeline) 127 | if err != nil { 128 | return nil, err 129 | } 130 | var results []models.User 131 | if err = cursor.All(ctx, &results); err != nil { 132 | return nil, err 133 | } 134 | return results, nil 135 | } 136 | -------------------------------------------------------------------------------- /common/database/operations/ugoira.go: -------------------------------------------------------------------------------- 1 | package operations 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/ShugetsuSoft/pixivel-back/common/models" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 9 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 10 | "go.mongodb.org/mongo-driver/bson" 11 | "go.mongodb.org/mongo-driver/mongo" 12 | ) 13 | 14 | func (ops *DatabaseOperations) InsertUgoira(ctx context.Context, ugoira *models.Ugoira) error { 15 | var err error 16 | is, err := ops.Flt.Exists(config.UgoiraTableName, utils.Itoa(ugoira.ID)) 17 | if err != nil { 18 | return err 19 | } 20 | ugoira.UpdateTime = time.Now() 21 | 22 | if is { 23 | goto REPLACE 24 | } else { 25 | _, err = ops.Cols.Ugoira.InsertOne(ctx, ugoira) 26 | 27 | if mongo.IsDuplicateKeyError(err) { 28 | _, err = ops.Flt.Add(config.UgoiraTableName, utils.Itoa(ugoira.ID)) 29 | if err != nil { 30 | return err 31 | } 32 | goto REPLACE 33 | } 34 | 35 | if err != nil { 36 | return err 37 | } 38 | 39 | _, err = ops.Flt.Add(config.UgoiraTableName, utils.Itoa(ugoira.ID)) 40 | if err != nil { 41 | return err 42 | } 43 | } 44 | 45 | return nil 46 | 47 | REPLACE: 48 | result, err := ops.Cols.Ugoira.ReplaceOne(ctx, bson.M{"_id": ugoira.ID}, ugoira) 49 | if err != nil { 50 | return err 51 | } 52 | if result.MatchedCount == 0 { 53 | _, err = ops.Cols.Ugoira.InsertOne(ctx, ugoira) 54 | if err != nil { 55 | return err 56 | } 57 | _, err = ops.Flt.Add(config.UgoiraTableName, utils.Itoa(ugoira.ID)) 58 | } 59 | 60 | return err 61 | } 62 | 63 | func (ops *DatabaseOperations) QueryUgoira(ctx context.Context, ugoiraId uint64) (*models.Ugoira, error) { 64 | is, err := ops.Flt.Exists(config.UgoiraTableName, utils.Itoa(ugoiraId)) 65 | 66 | if err != nil { 67 | return nil, err 68 | } 69 | 70 | if is { 71 | result := models.Ugoira{ 72 | Frames: []models.UgoiraFrame{}, 73 | } 74 | query := bson.M{"_id": ugoiraId} 75 | err := ops.Cols.Ugoira.FindOne(ctx, query).Decode(&result) 76 | if err != nil { 77 | if err == mongo.ErrNoDocuments { 78 | return nil, nil 79 | } else { 80 | return nil, err 81 | } 82 | } 83 | 84 | return &result, err 85 | } 86 | return nil, nil 87 | } 88 | 89 | func (ops *DatabaseOperations) IsUgoiraExist(ugoiraId uint64) (bool, error) { 90 | is, err := ops.Flt.Exists(config.UgoiraTableName, utils.Itoa(ugoiraId)) 91 | if err != nil { 92 | return false, err 93 | } 94 | return is, nil 95 | } 96 | 97 | func (ops *DatabaseOperations) DeleteUgoira(ctx context.Context, ugoiraId uint64) error { 98 | is, err := ops.Flt.Exists(config.UgoiraTableName, utils.Itoa(ugoiraId)) 99 | 100 | if err != nil { 101 | return err 102 | } 103 | 104 | if is { 105 | _, err := ops.Cols.Ugoira.DeleteOne(ctx, bson.M{"_id": ugoiraId}) 106 | if err != nil { 107 | if err == mongo.ErrNoDocuments { 108 | return nil 109 | } 110 | return err 111 | } 112 | 113 | return err 114 | } 115 | 116 | return nil 117 | } 118 | -------------------------------------------------------------------------------- /common/database/operations/user.go: -------------------------------------------------------------------------------- 1 | package operations 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/ShugetsuSoft/pixivel-back/common/convert" 8 | "github.com/ShugetsuSoft/pixivel-back/common/models" 9 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 10 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 11 | "github.com/olivere/elastic/v7" 12 | "go.mongodb.org/mongo-driver/bson" 13 | "go.mongodb.org/mongo-driver/mongo" 14 | ) 15 | 16 | func (ops *DatabaseOperations) InsertUser(ctx context.Context, user *models.User) error { 17 | var err error 18 | is, err := ops.Flt.Exists(config.UserTableName, utils.Itoa(user.ID)) 19 | if err != nil { 20 | return err 21 | } 22 | user.UpdateTime = time.Now() 23 | 24 | if is { 25 | goto REPLACE 26 | } else { 27 | user.IllustsUpdateTime = time.Unix(0, 0) 28 | user.IllustsCount = 0 29 | _, err = ops.Cols.User.InsertOne(ctx, user) 30 | 31 | if mongo.IsDuplicateKeyError(err) { 32 | _, err = ops.Flt.Add(config.UserTableName, utils.Itoa(user.ID)) 33 | if err != nil { 34 | return err 35 | } 36 | goto REPLACE 37 | } 38 | 39 | if err != nil { 40 | return err 41 | } 42 | 43 | _, err = ops.Flt.Add(config.UserTableName, utils.Itoa(user.ID)) 44 | if err != nil { 45 | return err 46 | } 47 | } 48 | 49 | return nil 50 | 51 | REPLACE: 52 | result, err := ops.Cols.User.ReplaceOne(ctx, bson.M{"_id": user.ID}, user) 53 | if err != nil { 54 | return err 55 | } 56 | if result.MatchedCount == 0 { 57 | _, err = ops.Cols.User.InsertOne(ctx, user) 58 | if err != nil { 59 | return err 60 | } 61 | _, err = ops.Flt.Add(config.UserTableName, utils.Itoa(user.ID)) 62 | } 63 | 64 | return err 65 | } 66 | 67 | func (ops *DatabaseOperations) UpdateUserIllustsTime(ctx context.Context, userId uint64) error { 68 | var err error 69 | is, err := ops.Flt.Exists(config.UserTableName, utils.Itoa(userId)) 70 | if err != nil { 71 | return err 72 | } 73 | if is { 74 | _, err = ops.Cols.User.UpdateOne(ctx, bson.M{"_id": userId}, bson.M{"$set": bson.M{ 75 | "illusts_update_time": time.Now(), 76 | "update_time": time.Now(), 77 | }}) 78 | return err 79 | } 80 | return nil 81 | } 82 | 83 | func (ops *DatabaseOperations) InsertUserSearch(ctx context.Context, user *models.User) error { 84 | usersearch := convert.User2UserSearch(user) 85 | err := ops.Sc.es.InsertDocument(ctx, config.UserSearchIndexName, utils.Itoa(user.ID), usersearch) 86 | if err != nil { 87 | return err 88 | } 89 | return nil 90 | } 91 | 92 | func (ops *DatabaseOperations) QueryUser(ctx context.Context, userId uint64, resultbanned bool) (*models.User, error) { 93 | is, err := ops.Flt.Exists(config.UserTableName, utils.Itoa(userId)) 94 | 95 | if err != nil { 96 | return nil, err 97 | } 98 | 99 | if is { 100 | result := models.User{ 101 | Image: models.UserImage{}, 102 | } 103 | err := ops.Cols.User.FindOne(ctx, bson.M{"_id": userId}).Decode(&result) 104 | 105 | if err != nil { 106 | if err == mongo.ErrNoDocuments { 107 | return nil, nil 108 | } else { 109 | return nil, err 110 | } 111 | } 112 | 113 | if result.Banned && !resultbanned { 114 | return nil, models.ErrorItemBanned 115 | } 116 | 117 | return &result, err 118 | } 119 | return nil, nil 120 | } 121 | 122 | func (ops *DatabaseOperations) QueryUsers(ctx context.Context, userIds []uint64, resultbanned bool) ([]models.User, error) { 123 | query := bson.M{"_id": bson.M{"$in": userIds}} 124 | cursor, err := ops.Cols.User.Find(ctx, query) 125 | defer cursor.Close(ctx) 126 | if err != nil { 127 | if err == mongo.ErrNoDocuments { 128 | return nil, nil 129 | } else { 130 | return nil, err 131 | } 132 | } 133 | 134 | users := make([]models.User, 0, len(userIds)) 135 | 136 | for cursor.Next(ctx) { 137 | result := models.User{ 138 | Image: models.UserImage{}, 139 | } 140 | err := cursor.Decode(&result) 141 | if err != nil { 142 | return nil, err 143 | } 144 | 145 | if result.Banned && !resultbanned { 146 | continue 147 | } 148 | 149 | users = append(users, result) 150 | } 151 | 152 | return users, err 153 | } 154 | 155 | func (ops *DatabaseOperations) SearchUserSuggest(ctx context.Context, keyword string) ([]string, error) { 156 | source := elastic.NewSearchSource(). 157 | Suggester(ops.Sc.es.Suggest("name-completion-suggest").Field("name.suggest").Text(keyword).Fuzziness(2).Analyzer("kuromoji")). 158 | FetchSource(false).TrackScores(true) 159 | query := ops.Sc.es.Search(config.UserSearchIndexName).Source(nil). 160 | SearchSource(source) 161 | 162 | results, err := ops.Sc.es.DoSearch(ctx, query) 163 | if err != nil { 164 | return nil, err 165 | } 166 | suggests := results.Suggest["name-completion-suggest"][0].Options 167 | res := make([]string, len(suggests)) 168 | for i, suggest := range suggests { 169 | res[i] = suggest.Text 170 | } 171 | return res, nil 172 | } 173 | 174 | func (ops *DatabaseOperations) SearchUser(ctx context.Context, keyword string, page int, limit int, resultbanned bool) ([]models.User, int64, []float64, []*string, error) { 175 | query := ops.Sc.es.Search(config.UserSearchIndexName). 176 | Query(ops.Sc.es.BoolQuery(). 177 | Should(ops.Sc.es.Query("name", keyword).Analyzer("kuromoji").Boost(3)). 178 | Should(elastic.NewMatchQuery("name.keyword", keyword).Boost(2)). 179 | Should(ops.Sc.es.Query("bio", keyword).Analyzer("kuromoji").Boost(1)), 180 | ). 181 | Size(limit).From(page * limit). 182 | Highlight(elastic.NewHighlight().Field("name")). 183 | FetchSourceContext(elastic.NewFetchSourceContext(true).Include("_id")).TrackScores(true) 184 | 185 | query = query.Sort("_score", false) 186 | 187 | results, err := ops.Sc.es.DoSearch(ctx, query) 188 | if err != nil { 189 | return nil, 0, nil, nil, err 190 | } 191 | 192 | if results.Hits.TotalHits.Value > 0 { 193 | scores := make([]float64, 0, len(results.Hits.Hits)) 194 | highlights := make([]*string, 0, len(results.Hits.Hits)) 195 | userids := make([]uint64, len(results.Hits.Hits)) 196 | for i, hit := range results.Hits.Hits { 197 | userids[i] = utils.Atoi(hit.Id) 198 | if hit.Score != nil { 199 | scores = append(scores, *hit.Score) 200 | } else { 201 | scores = append(scores, -1) 202 | } 203 | if highl, ok := hit.Highlight["name"]; ok { 204 | highlights = append(highlights, &highl[0]) 205 | } else { 206 | highlights = append(highlights, nil) 207 | } 208 | } 209 | 210 | users, err := ops.QueryUsers(ctx, userids, resultbanned) 211 | if err != nil { 212 | return nil, 0, nil, nil, err 213 | } 214 | 215 | usersmap := make(map[uint64]models.User) 216 | for _, user := range users { 217 | usersmap[user.ID] = user 218 | } 219 | 220 | result := make([]models.User, 0, len(results.Hits.Hits)) 221 | for _, userid := range userids { 222 | if _, exist := usersmap[userid]; exist { 223 | result = append(result, usersmap[userid]) 224 | } 225 | } 226 | 227 | return result, results.Hits.TotalHits.Value, scores, highlights, nil 228 | } else { 229 | return nil, 0, nil, nil, models.ErrorNoResult 230 | } 231 | return nil, 0, nil, nil, err 232 | } 233 | 234 | func (ops *DatabaseOperations) DeleteUser(ctx context.Context, userId uint64) error { 235 | _, err := ops.Cols.User.DeleteOne(ctx, bson.M{"_id": userId}) 236 | if err != nil { 237 | if err == mongo.ErrNoDocuments { 238 | return nil 239 | } 240 | return err 241 | } 242 | return err 243 | } 244 | 245 | func (ops *DatabaseOperations) ClearUserIllusts(ctx context.Context, userId uint64) error { 246 | illusts, err := ops.QueryIllustByUser(ctx, userId, true) 247 | if err != nil { 248 | return err 249 | } 250 | 251 | for i := 0; i < len(illusts); i++ { 252 | err = ops.DeleteIllust(ctx, (illusts)[i].ID) 253 | if err != nil { 254 | return err 255 | } 256 | } 257 | 258 | is, err := ops.Flt.Exists(config.UserTableName, utils.Itoa(userId)) 259 | 260 | if err != nil { 261 | return err 262 | } 263 | 264 | if is { 265 | _, err = ops.Cols.User.UpdateOne(ctx, bson.M{"_id": userId}, bson.M{"$set": bson.M{ 266 | "illusts_update_time": time.Unix(0, 0), 267 | "illusts_count": 0, 268 | "update_time": time.Now(), 269 | }}) 270 | return err 271 | } 272 | return nil 273 | } 274 | 275 | func (ops *DatabaseOperations) SetIllustsCount(ctx context.Context, userId uint64, count uint) error { 276 | var err error 277 | is, err := ops.Flt.Exists(config.UserTableName, utils.Itoa(userId)) 278 | if err != nil { 279 | return err 280 | } 281 | if is { 282 | _, err = ops.Cols.User.UpdateOne(ctx, bson.M{"_id": userId}, bson.M{"$set": bson.M{ 283 | "illusts_update_time": time.Now(), 284 | "illusts_count": count, 285 | "update_time": time.Now(), 286 | }}) 287 | return err 288 | } 289 | return nil 290 | } 291 | 292 | func (ops *DatabaseOperations) IsUserExist(userId uint64) (bool, error) { 293 | is, err := ops.Flt.Exists(config.UserTableName, utils.Itoa(userId)) 294 | if err != nil { 295 | return false, err 296 | } 297 | return is, nil 298 | } 299 | -------------------------------------------------------------------------------- /common/database/tasktracer/tracer.go: -------------------------------------------------------------------------------- 1 | package tasktracer 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "log" 7 | "strings" 8 | "sync" 9 | "time" 10 | 11 | "github.com/ShugetsuSoft/pixivel-back/common/database/drivers" 12 | "github.com/ShugetsuSoft/pixivel-back/common/models" 13 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 14 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 15 | ) 16 | 17 | type TaskTracer struct { 18 | redis *drivers.RedisPool 19 | tracerchan string 20 | } 21 | 22 | type TaskListener struct { 23 | cond *sync.Cond 24 | finish string 25 | cancle context.CancelFunc 26 | redis *drivers.RedisClient 27 | } 28 | 29 | func NewTaskTracer(redis *drivers.RedisPool, tracerchan string) *TaskTracer { 30 | return &TaskTracer{ 31 | redis, 32 | tracerchan, 33 | } 34 | } 35 | 36 | func (tc *TaskTracer) NewTaskGroup(params map[string]string, tasknum uint) (string, bool, error) { 37 | tid := utils.HashMap(params) 38 | cli := tc.redis.NewRedisClient() 39 | defer cli.Close() 40 | isexist, err := cli.KeyExist(tid) 41 | if err != nil { 42 | return "", false, err 43 | } 44 | if isexist { 45 | return tid, true, nil 46 | } 47 | err = cli.SetValueExpire(tid, utils.Itoa(tasknum), 60*10) 48 | if err != nil { 49 | return "", false, err 50 | } 51 | return tid, false, nil 52 | } 53 | 54 | func (tc *TaskTracer) NewTaskByNum(tid string, tasknum uint) error { 55 | if tid == "" { 56 | return nil 57 | } 58 | cli := tc.redis.NewRedisClient() 59 | defer cli.Close() 60 | _, err := cli.IncreaseBy(tid, tasknum) 61 | return err 62 | } 63 | 64 | func (tc *TaskTracer) NewTask(tid string) error { 65 | if tid == "" { 66 | return nil 67 | } 68 | cli := tc.redis.NewRedisClient() 69 | defer cli.Close() 70 | _, err := cli.Increase(tid) 71 | return err 72 | } 73 | 74 | func (tc *TaskTracer) RemoveTaskGroup(tid string) error { 75 | if tid == "" { 76 | return nil 77 | } 78 | cli := tc.redis.NewRedisClient() 79 | defer cli.Close() 80 | return cli.DeleteValue(tid) 81 | } 82 | 83 | func (tc *TaskTracer) FinishTask(tid string) bool { 84 | if tid == "" { 85 | return true 86 | } 87 | cli := tc.redis.NewRedisClient() 88 | defer cli.Close() 89 | num, err := cli.Decrease(tid) 90 | if err != nil { 91 | telemetry.Log(telemetry.Label{"pos": "tasktracer"}, err.Error()) 92 | log.Fatal(err) 93 | } 94 | if num <= 0 { 95 | err = tc.RemoveTaskGroup(tid) 96 | if err != nil { 97 | telemetry.Log(telemetry.Label{"pos": "tasktracer"}, err.Error()) 98 | log.Fatal(err) 99 | } 100 | err = cli.Publish(tc.tracerchan, tid) 101 | if err != nil { 102 | telemetry.Log(telemetry.Label{"pos": "tasktracer"}, err.Error()) 103 | log.Fatal(err) 104 | } 105 | return true 106 | } 107 | return false 108 | } 109 | 110 | func (tc *TaskTracer) FailTask(tid string, info string) bool { 111 | if tid == "" { 112 | return true 113 | } 114 | err := tc.RemoveTaskGroup(tid) 115 | if err != nil { 116 | telemetry.Log(telemetry.Label{"pos": "tasktracer"}, err.Error()) 117 | log.Fatal(err) 118 | } 119 | cli := tc.redis.NewRedisClient() 120 | defer cli.Close() 121 | err = cli.Publish(tc.tracerchan, tid+"-ERR-"+info) 122 | if err != nil { 123 | telemetry.Log(telemetry.Label{"pos": "tasktracer"}, err.Error()) 124 | log.Fatal(err) 125 | } 126 | return true 127 | } 128 | 129 | func (tc *TaskTracer) NewListener(ctx context.Context) *TaskListener { 130 | ctx, cancel := context.WithCancel(ctx) 131 | cond := sync.NewCond(&sync.Mutex{}) 132 | cli := tc.redis.NewRedisClient() 133 | tracerchan, err := cli.Subscribe(ctx, tc.tracerchan) 134 | if err != nil { 135 | telemetry.Log(telemetry.Label{"pos": "tasktracer"}, err.Error()) 136 | log.Fatal(err) 137 | } 138 | tasklistener := &TaskListener{ 139 | cond, 140 | "", 141 | cancel, 142 | cli, 143 | } 144 | go func() { 145 | for { 146 | tid, ok := <-tracerchan 147 | if !ok { 148 | return 149 | } 150 | tasklistener.cond.L.Lock() 151 | tasklistener.finish = tid 152 | tasklistener.cond.L.Unlock() 153 | tasklistener.cond.Broadcast() 154 | } 155 | }() 156 | return tasklistener 157 | } 158 | 159 | func (tl *TaskListener) CloseChan() { 160 | tl.cancle() 161 | tl.redis.Close() 162 | } 163 | 164 | func (tl *TaskListener) WaitFor(ctx context.Context, tid string, delay time.Duration) error { 165 | timeout := time.After(delay) 166 | signal := make(chan error, 1) 167 | go func() { 168 | tl.cond.L.Lock() 169 | for { 170 | if tl.finish == tid { 171 | tl.cond.L.Unlock() 172 | signal <- nil 173 | return 174 | } else if strings.Contains(tl.finish, tid) { 175 | tl.cond.L.Unlock() 176 | signal <- errors.New(strings.Split(tl.finish, "-")[2]) 177 | return 178 | } 179 | tl.cond.Wait() 180 | } 181 | }() 182 | select { 183 | case err := <-signal: 184 | return err 185 | case <-timeout: 186 | return models.ErrorTimeOut 187 | case <-ctx.Done(): 188 | return models.ErrorTimeOut 189 | } 190 | } 191 | -------------------------------------------------------------------------------- /common/models/config.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | type Modes int 4 | 5 | const ( 6 | NormalMode Modes = iota 7 | ArchiveMode 8 | ) 9 | 10 | type Config struct { 11 | Mongodb MongoDBConfig `yaml:"mongodb"` 12 | Rabbitmq RabbitMQConfig `yaml:"rabbitmq"` 13 | Elasticsearch ElasticSearchConfig `yaml:"elasticsearch"` 14 | Redis RedisConfig `yaml:"redis"` 15 | Neardb NearDBConfig `yaml:"neardb"` 16 | Responser ResponserConfig `yaml:"responser"` 17 | General GeneralConfig `yaml:"general"` 18 | Spider SpiderConfig `yaml:"spider"` 19 | } 20 | 21 | type ResponserConfig struct { 22 | Listen string `yaml:"listen"` 23 | Debug bool `yaml:"debug"` 24 | Mode Modes `yaml:"mode"` 25 | EnableForceFetch bool `yaml:"enable-force-fetch"` 26 | } 27 | 28 | type SpiderConfig struct { 29 | PixivToken string `yaml:"pixiv-token"` 30 | CrawlingThreads int `yaml:"crawling-threads"` 31 | } 32 | 33 | type GeneralConfig struct { 34 | SpiderRetry uint `yaml:"spider-retry"` 35 | Prometheus string `yaml:"prometheus-listen"` 36 | Loki string `yaml:"loki-uri"` 37 | } 38 | 39 | type MongoDBConfig struct { 40 | URI string `yaml:"uri"` 41 | } 42 | 43 | type RabbitMQConfig struct { 44 | URI string `yaml:"uri"` 45 | } 46 | 47 | type ElasticSearchConfig struct { 48 | URI string `yaml:"uri"` 49 | User string `yaml:"user"` 50 | Pass string `yaml:"password"` 51 | } 52 | 53 | type RedisConfig struct { 54 | CacheRedis CacheRedisConfig `yaml:"cache"` 55 | MessageRedis MessageRedisConfig `yaml:"message"` 56 | } 57 | 58 | type CacheRedisConfig struct { 59 | URI string `yaml:"uri"` 60 | } 61 | 62 | type MessageRedisConfig struct { 63 | URI string `yaml:"uri"` 64 | } 65 | 66 | type NearDBConfig struct { 67 | URI string `yaml:"uri"` 68 | } 69 | -------------------------------------------------------------------------------- /common/models/database.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import "time" 4 | 5 | type UserImage struct { 6 | Url string `bson:"url"` 7 | BigUrl string `bson:"bigUrl"` 8 | Background string `bson:"background"` 9 | } 10 | 11 | type User struct { 12 | ID uint64 `bson:"_id"` 13 | UpdateTime time.Time `bson:"update_time"` 14 | 15 | Name string `bson:"name"` 16 | Bio string `bson:"bio,omitempty"` 17 | Image UserImage `bson:"image"` 18 | IllustsUpdateTime time.Time `bson:"illusts_update_time"` 19 | IllustsCount uint `bson:"illusts_count"` 20 | 21 | Banned bool `bson:"banned"` 22 | } 23 | 24 | type IllustTag struct { 25 | Name string `bson:"name" json:"name"` 26 | Translation string `bson:"translation,omitempty" json:"translation,omitempty"` 27 | } 28 | 29 | type IllustStatistic struct { 30 | Bookmarks uint `bson:"bookmarks" json:"bookmarks"` 31 | Likes uint `bson:"likes" json:"likes"` 32 | Comments uint `bson:"comments" json:"comments"` 33 | Views uint `bson:"views" json:"views"` 34 | } 35 | 36 | type Illust struct { 37 | ID uint64 `bson:"_id"` 38 | UpdateTime time.Time `bson:"update_time"` 39 | 40 | Title string `bson:"title"` 41 | AltTitle string `bson:"altTitle"` 42 | Description string `bson:"description,omitempty"` 43 | Type uint `bson:"type"` 44 | CreateDate time.Time `bson:"createDate"` 45 | UploadDate time.Time `bson:"uploadDate"` 46 | Sanity uint `bson:"sanity"` 47 | Width uint `bson:"width"` 48 | Height uint `bson:"height"` 49 | PageCount uint `bson:"pageCount"` 50 | Tags []IllustTag `bson:"tags"` 51 | Statistic IllustStatistic `bson:"statistic"` 52 | User uint `bson:"user"` 53 | Image time.Time `bson:"image"` 54 | AIType uint `bson:"aiType,omitempty"` 55 | 56 | Popularity uint `bson:"popularity"` 57 | Banned bool `bson:"banned"` 58 | } 59 | 60 | type RankIllust struct { 61 | ID uint64 `bson:"id"` 62 | Rank uint `bson:"rank"` 63 | } 64 | 65 | type RankAggregateResult struct { 66 | Content Illust `bson:"content"` 67 | } 68 | 69 | type Rank struct { 70 | Date string `bson:"date"` 71 | Mode string `bson:"mode"` 72 | Content string `bson:"content"` 73 | Illusts []RankIllust `bson:"illusts"` 74 | } 75 | 76 | type Ugoira struct { 77 | ID uint64 `bson:"_id"` 78 | UpdateTime time.Time `bson:"update_time"` 79 | 80 | Image time.Time `bson:"image"` 81 | MimeType string `bson:"mimeType"` 82 | Frames []UgoiraFrame `bson:"frames"` 83 | } 84 | 85 | type UgoiraFrame struct { 86 | File string `bson:"file"` 87 | Delay int `bson:"delay"` 88 | } 89 | 90 | type IllustSearch struct { 91 | Title string `json:"title"` 92 | AltTitle string `json:"alt_title,omitempty"` 93 | Description string `json:"description,omitempty"` 94 | Type uint `json:"type"` 95 | CreateDate time.Time `json:"create_date"` 96 | Sanity uint `json:"sanity"` 97 | Popularity uint `json:"popularity"` 98 | User uint `json:"user"` 99 | Tags []IllustTag `json:"tags"` 100 | } 101 | 102 | type UserSearch struct { 103 | Name string `json:"name"` 104 | Bio string `json:"bio,omitempty"` 105 | } 106 | 107 | const ( 108 | IllustSearchMapping = `{"mappings":{"properties":{"alt_title":{"type":"text","analyzer":"kuromoji","search_analyzer":"kuromoji","fields":{"keyword":{"type":"keyword","ignore_above":256},"suggest":{"type": "completion","analyzer":"kuromoji"}}},"create_date":{"type":"date"},"description":{"type":"text","analyzer":"kuromoji","search_analyzer":"kuromoji"},"sanity":{"type":"short"},"popularity":{"type":"long"},"title":{"type":"text","analyzer":"kuromoji","search_analyzer":"kuromoji","fields":{"keyword":{"type":"keyword","ignore_above":256},"suggest":{"type": "completion","analyzer":"kuromoji"}}},"type":{"type":"short"},"user":{"type":"long"},"tags":{"type":"nested", "properties":{"name":{"type":"keyword","fields":{"suggest":{"type": "completion","analyzer":"kuromoji"},"fuzzy":{"type":"text","analyzer":"kuromoji","search_analyzer":"kuromoji"}}},"translation":{"type":"keyword","fields":{"suggest":{"type": "completion","analyzer":"smartcn"},"fuzzy":{"type":"text","analyzer":"smartcn","search_analyzer":"smartcn"}}}}}}}}` 109 | UserSearchMapping = `{"mappings":{"properties":{"bio":{"type":"text","analyzer":"kuromoji","search_analyzer":"kuromoji","fields":{"keyword":{"type":"keyword","ignore_above":256}}},"name":{"type":"text","analyzer":"kuromoji","search_analyzer":"kuromoji","fields":{"keyword":{"type":"keyword","ignore_above":256},"suggest":{"type": "completion","analyzer":"kuromoji"}}}}}}` 110 | ) 111 | -------------------------------------------------------------------------------- /common/models/errors.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import "errors" 4 | 5 | var ( 6 | ErrorRetrivingFinishedTask = errors.New("error In Retryving Finished Task") 7 | ErrorIndexExist = errors.New("error Index Already Existed") 8 | ErrorItemBanned = errors.New("error Item Banned") 9 | ErrorNoResult = errors.New("error No Result") 10 | ErrorChannelClosed = errors.New("channel closed") 11 | ErrorTimeOut = errors.New("time Out") 12 | ErrorArchiveMode = errors.New("in Archive Mode") 13 | 14 | InternalErrorLoginNeeded = errors.New("login needed") 15 | ) 16 | -------------------------------------------------------------------------------- /common/models/interface.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | type MessageQueue interface { 4 | Publish(string, MQMessage) error 5 | PublishToExchange(string, MQMessage) error 6 | Consume(string) (<-chan MQMessage, error) 7 | Ack(uint64) error 8 | Reject(uint64) error 9 | } 10 | 11 | type Filter interface { 12 | Add(string, string) (bool, error) 13 | Exists(string, string) (bool, error) 14 | } 15 | -------------------------------------------------------------------------------- /common/models/message.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | const ( 4 | CrawlIllustDetail uint = iota 5 | CrawlUserDetail 6 | CrawlUserIllusts 7 | CrawlRankIllusts 8 | CrawlUgoiraDetail 9 | 10 | CrawlError 11 | ) 12 | 13 | type IDList []uint64 14 | 15 | type IDWithPos struct { 16 | ID uint64 17 | Pos uint 18 | } 19 | 20 | type UserIllustsResponse struct { 21 | UserID uint64 22 | Illusts IDList 23 | } 24 | 25 | type RankIllustsResponseMessage struct { 26 | Mode string 27 | Date string 28 | Content string 29 | Page int 30 | Next bool 31 | Illusts []IDWithPos 32 | } 33 | 34 | type MQMessage struct { 35 | Data []byte 36 | Tag uint64 37 | Priority uint8 38 | } 39 | 40 | type CrawlTask struct { 41 | Group string 42 | Type uint 43 | Params map[string]string 44 | RetryCount uint 45 | } 46 | 47 | type CrawlResponse struct { 48 | Group string 49 | Type uint 50 | Response interface{} 51 | } 52 | 53 | type CrawlErrorResponse struct { 54 | TaskType uint 55 | Params map[string]string 56 | Message string 57 | } 58 | -------------------------------------------------------------------------------- /common/models/pixiv.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import "time" 4 | 5 | type IllustRawUrls struct { 6 | Mini string `json:"mini"` 7 | Thumb string `json:"thumb"` 8 | Small string `json:"small"` 9 | Regular string `json:"regular"` 10 | Original string `json:"original"` 11 | } 12 | type IllustRawTranslation struct { 13 | En string `json:"en"` 14 | } 15 | type IllustRawTag struct { 16 | Tag string `json:"tag"` 17 | Translation IllustRawTranslation `json:"translation"` 18 | } 19 | 20 | type IllustRawTagPre struct { 21 | Tags []IllustRawTag `json:"tags"` 22 | } 23 | 24 | type IllustRaw struct { 25 | ID uint64 `json:"id,string"` 26 | Title string `json:"title"` 27 | Description string `json:"description"` 28 | IllustType uint `json:"illustType"` 29 | CreateDate time.Time `json:"createDate"` 30 | UploadDate time.Time `json:"uploadDate"` 31 | Restrict uint `json:"restrict"` 32 | XRestrict uint `json:"xRestrict"` 33 | Sl uint `json:"sl"` 34 | Urls IllustRawUrls `json:"urls"` 35 | Tags IllustRawTagPre `json:"tags"` 36 | Alt string `json:"alt"` 37 | UserID uint `json:"userId,string"` 38 | AIType uint `json:"aiType"` 39 | 40 | Width uint `json:"width"` 41 | Height uint `json:"height"` 42 | PageCount uint `json:"pageCount"` 43 | 44 | BookmarkCount uint `json:"bookmarkCount"` 45 | LikeCount uint `json:"likeCount"` 46 | CommentCount uint `json:"commentCount"` 47 | ViewCount uint `json:"viewCount"` 48 | } 49 | 50 | type IllustRawResponse struct { 51 | Error bool `json:"error"` 52 | Message string `json:"message"` 53 | Body IllustRaw `json:"body"` 54 | } 55 | 56 | type UserRawBackground struct { 57 | URL string `json:"url"` 58 | } 59 | 60 | type UserRaw struct { 61 | UserID uint64 `json:"userId,string"` 62 | Name string `json:"name"` 63 | Image string `json:"image"` 64 | ImageBig string `json:"imageBig"` 65 | Comment string `json:"comment"` 66 | Background UserRawBackground `json:"background"` 67 | } 68 | 69 | type UserRawResponse struct { 70 | Error bool `json:"error"` 71 | Message string `json:"message"` 72 | Body UserRaw `json:"body"` 73 | } 74 | 75 | type UserIllustsRaw struct { 76 | Illusts interface{} `json:"illusts"` 77 | Manga interface{} `json:"manga"` 78 | } 79 | 80 | type UserIllustsRawResponse struct { 81 | Error bool `json:"error"` 82 | Message string `json:"message"` 83 | Body UserIllustsRaw `json:"body"` 84 | } 85 | 86 | type RankIllustsRawResponse struct { 87 | Contents []RankIllustRaw `json:"contents"` 88 | Mode string `json:"mode"` 89 | Page int `json:"page"` 90 | Date string `json:"date"` 91 | Next interface{} `json:"next"` 92 | Content string `json:"content"` 93 | } 94 | 95 | type RankIllustRaw struct { 96 | IllustID uint64 `json:"illust_id"` 97 | UserID uint `json:"user_id"` 98 | Rank uint `json:"rank"` 99 | YesRank uint `json:"yes_rank"` 100 | RatingCount uint `json:"rating_count"` 101 | IllustUploadTimestamp uint `json:"illust_upload_timestamp"` 102 | } 103 | 104 | type UgoiraRawResponse struct { 105 | Error bool `json:"error"` 106 | Message string `json:"message"` 107 | Body UgoiraRaw `json:"body"` 108 | } 109 | 110 | type UgoiraFrameRaw struct { 111 | File string `json:"file"` 112 | Delay int `json:"delay"` 113 | } 114 | 115 | type UgoiraRaw struct { 116 | Src string `json:"src"` 117 | Originalsrc string `json:"originalSrc"` 118 | MimeType string `json:"mime_type"` 119 | Frames []UgoiraFrameRaw `json:"frames"` 120 | } 121 | 122 | type ErrorRawResponse struct { 123 | Message string `json:"message"` 124 | } 125 | -------------------------------------------------------------------------------- /common/models/response.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | type UserImageResponse struct { 4 | Url string `json:"url"` 5 | BigUrl string `json:"bigUrl"` 6 | Background string `json:"background,omitempty"` 7 | } 8 | 9 | type UserResponse struct { 10 | ID uint64 `json:"id"` 11 | Name string `json:"name"` 12 | Bio string `json:"bio"` 13 | Image UserImageResponse `json:"image"` 14 | } 15 | 16 | type IllustTagResponse struct { 17 | Name string `json:"name"` 18 | Translation string `json:"translation,omitempty"` 19 | } 20 | 21 | type IllustStatisticResponse struct { 22 | Bookmarks uint `json:"bookmarks"` 23 | Likes uint `json:"likes"` 24 | Comments uint `json:"comments"` 25 | Views uint `json:"views"` 26 | } 27 | 28 | type IllustResponse struct { 29 | ID uint64 `json:"id"` 30 | Title string `json:"title"` 31 | AltTitle string `json:"altTitle"` 32 | Description string `json:"description"` 33 | Type uint `json:"type"` 34 | CreateDate string `json:"createDate"` 35 | UploadDate string `json:"uploadDate"` 36 | Sanity uint `json:"sanity"` 37 | Width uint `json:"width"` 38 | Height uint `json:"height"` 39 | PageCount uint `json:"pageCount"` 40 | Tags []IllustTagResponse `json:"tags"` 41 | Statistic IllustStatisticResponse `json:"statistic"` 42 | User *UserResponse `json:"user,omitempty"` 43 | Image string `json:"image"` 44 | AIType uint `json:"aiType"` 45 | } 46 | 47 | type IllustsResponse struct { 48 | Illusts []IllustResponse `json:"illusts"` 49 | HasNext bool `json:"has_next"` 50 | } 51 | 52 | type UsersResponse struct { 53 | Users []UserResponse `json:"users"` 54 | HasNext bool `json:"has_next"` 55 | } 56 | 57 | type UsersSearchResponse struct { 58 | Users []UserResponse `json:"users"` 59 | Scores []float64 `json:"scores"` 60 | HighLight []*string `json:"highlight,omitempty"` 61 | HasNext bool `json:"has_next"` 62 | } 63 | 64 | type IllustsSearchResponse struct { 65 | Illusts []IllustResponse `json:"illusts"` 66 | Scores []float64 `json:"scores,omitempty"` 67 | HighLight []*string `json:"highlight,omitempty"` 68 | HasNext bool `json:"has_next"` 69 | } 70 | 71 | type SearchSuggestResponse struct { 72 | SuggestWords []string `json:"suggest_words"` 73 | } 74 | 75 | type SearchSuggestTagsResponse struct { 76 | SuggestTags []IllustTagResponse `json:"suggest_tags"` 77 | } 78 | 79 | type UgoiraResponse struct { 80 | ID uint64 `json:"id"` 81 | Image string `json:"image"` 82 | MimeType string `json:"mime_type"` 83 | Frames []UgoiraFrameResponse `json:"frames"` 84 | } 85 | 86 | type UgoiraFrameResponse struct { 87 | File string `json:"file"` 88 | Delay int `json:"delay"` 89 | } 90 | -------------------------------------------------------------------------------- /common/utils/config/const.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | const ( 4 | DatabaseName = "Pixivel" 5 | IllustTableName = "Illust" 6 | UserTableName = "User" 7 | RankTableName = "Rank" 8 | UgoiraTableName = "Ugoira" 9 | 10 | IllustSearchIndexName = "illust" 11 | UserSearchIndexName = "user" 12 | 13 | TaskTracerChannel = "TaskTracer" 14 | CrawlTaskQueue = "CrawlTasks" 15 | CrawlResponsesMongodb = "CrawlResponsesMongodb" 16 | CrawlResponsesElastic = "CrawlResponsesElastic" 17 | CrawlResponsesExchange = "CrawlResponses" 18 | 19 | UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36" 20 | 21 | RecommendDrift = 0.00001 22 | ) 23 | -------------------------------------------------------------------------------- /common/utils/hash.go: -------------------------------------------------------------------------------- 1 | package utils 2 | 3 | import ( 4 | "crypto/md5" 5 | "encoding/hex" 6 | ) 7 | 8 | func HashMap(amap map[string]string) string { 9 | h := md5.New() 10 | 11 | for k, v := range amap { 12 | h.Write(StringIn(k + "=" + v + ";")) 13 | } 14 | return hex.EncodeToString(h.Sum(nil)) 15 | } 16 | -------------------------------------------------------------------------------- /common/utils/msgpack.go: -------------------------------------------------------------------------------- 1 | package utils 2 | 3 | import "github.com/vmihailenco/msgpack/v5" 4 | 5 | func MsgPack(data interface{}) ([]byte, error) { 6 | return msgpack.Marshal(data) 7 | } 8 | 9 | func MsgUnpack(raw []byte, data interface{}) error { 10 | return msgpack.Unmarshal(raw, data) 11 | } 12 | -------------------------------------------------------------------------------- /common/utils/string.go: -------------------------------------------------------------------------------- 1 | package utils 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/models" 5 | "strings" 6 | "unicode" 7 | ) 8 | 9 | func RemoveSpecialChars(in string) string { 10 | return strings.Map(func(r rune) rune { 11 | if unicode.IsPrint(r) { 12 | return r 13 | } 14 | return -1 15 | }, in) 16 | } 17 | 18 | func RemoveSpecialCharsTags(tags []models.IllustTag) []models.IllustTag { 19 | result := make([]models.IllustTag, len(tags)) 20 | for i := 0; i < len(tags); i++ { 21 | result[i].Name = RemoveSpecialChars(tags[i].Name) 22 | result[i].Translation = RemoveSpecialChars(tags[i].Translation) 23 | } 24 | return result 25 | } 26 | -------------------------------------------------------------------------------- /common/utils/telemetry/logger.go: -------------------------------------------------------------------------------- 1 | package telemetry 2 | 3 | import ( 4 | "log" 5 | "time" 6 | 7 | "github.com/ShugetsuSoft/loki-client-go/lib" 8 | ) 9 | 10 | type Label map[string]string 11 | 12 | func Log(label Label, s string) { 13 | log.Printf("[%+v] %s", label, s) 14 | label["type"] = LogType 15 | err := LoglokiIns.WriteLog(lib.Label(label), s, time.Now()) 16 | if err != nil { 17 | log.Println(err) 18 | } 19 | } 20 | -------------------------------------------------------------------------------- /common/utils/telemetry/loki.go: -------------------------------------------------------------------------------- 1 | package telemetry 2 | 3 | import ( 4 | loki "github.com/ShugetsuSoft/loki-client-go" 5 | ) 6 | 7 | var LoglokiIns = &loki.LokiClient{} 8 | var LogType = "" 9 | 10 | func RunLoki(uri string, logType string) chan error { 11 | LoglokiIns = loki.NewLokiClient(uri) 12 | LogType = logType 13 | return LoglokiIns.RunPush() 14 | } 15 | -------------------------------------------------------------------------------- /common/utils/telemetry/prometheus.go: -------------------------------------------------------------------------------- 1 | package telemetry 2 | 3 | import ( 4 | "github.com/prometheus/client_golang/prometheus" 5 | "github.com/prometheus/client_golang/prometheus/promhttp" 6 | "log" 7 | "net/http" 8 | ) 9 | 10 | var ( 11 | RequestsCount = prometheus.NewCounterVec(prometheus.CounterOpts{ 12 | Name: "total_requests_count", 13 | }, []string{"handler"}) 14 | RequestsHitCache = prometheus.NewCounter(prometheus.CounterOpts{ 15 | Name: "requests_cache_hit_count", 16 | }) 17 | RequestsErrorCount = prometheus.NewCounterVec(prometheus.CounterOpts{ 18 | Name: "requests_error_count", 19 | }, []string{"handler"}) 20 | 21 | SpiderTaskCount = prometheus.NewCounter(prometheus.CounterOpts{ 22 | Name: "spider_task_count", 23 | }) 24 | SpiderErrorTaskCount = prometheus.NewCounter(prometheus.CounterOpts{ 25 | Name: "spider_error_task_count", 26 | }) 27 | ) 28 | 29 | func RegisterResponser() { 30 | prometheus.MustRegister(RequestsCount) 31 | prometheus.MustRegister(RequestsHitCache) 32 | prometheus.MustRegister(RequestsErrorCount) 33 | } 34 | 35 | func RegisterSpider() { 36 | prometheus.MustRegister(SpiderTaskCount) 37 | prometheus.MustRegister(SpiderErrorTaskCount) 38 | } 39 | 40 | func RegisterStorer() { 41 | 42 | } 43 | 44 | func RunPrometheus(addr string) { 45 | http.Handle("/metrics", promhttp.Handler()) 46 | log.Fatal(http.ListenAndServe(addr, nil)) 47 | } 48 | -------------------------------------------------------------------------------- /common/utils/typeconv.go: -------------------------------------------------------------------------------- 1 | package utils 2 | 3 | import ( 4 | "bytes" 5 | "crypto/md5" 6 | "encoding/binary" 7 | "encoding/json" 8 | "fmt" 9 | "strconv" 10 | "unsafe" 11 | ) 12 | 13 | func HashStruct(item interface{}) string { 14 | jsonBytes, _ := json.Marshal(item) 15 | return fmt.Sprintf("%x", md5.Sum(jsonBytes)) 16 | } 17 | 18 | func StringOut(bye []byte) string { 19 | return *(*string)(unsafe.Pointer(&bye)) 20 | } 21 | 22 | func StringIn(strings string) []byte { 23 | x := (*[2]uintptr)(unsafe.Pointer(&strings)) 24 | return *(*[]byte)(unsafe.Pointer(&[3]uintptr{x[0], x[1], x[1]})) 25 | } 26 | 27 | func IntIn(n int) []byte { 28 | data := int64(n) 29 | bytebuf := bytes.NewBuffer([]byte{}) 30 | binary.Write(bytebuf, binary.BigEndian, data) 31 | return bytebuf.Bytes() 32 | } 33 | 34 | func IntOut(bye []byte) int { 35 | bytebuff := bytes.NewBuffer(bye) 36 | var data int64 37 | binary.Read(bytebuff, binary.BigEndian, &data) 38 | return int(data) 39 | } 40 | 41 | func UintIn(n uint) []byte { 42 | data := uint64(n) 43 | bytebuf := bytes.NewBuffer([]byte{}) 44 | binary.Write(bytebuf, binary.BigEndian, data) 45 | return bytebuf.Bytes() 46 | } 47 | 48 | func Uint64Out(bye []byte) uint64 { 49 | bytebuff := bytes.NewBuffer(bye) 50 | var data uint64 51 | binary.Read(bytebuff, binary.BigEndian, &data) 52 | return data 53 | } 54 | 55 | func UintOut(bye []byte) uint { 56 | return uint(Uint64Out(bye)) 57 | } 58 | 59 | func Itoa(a interface{}) string { 60 | if v, p := a.(int); p { 61 | return strconv.Itoa(v) 62 | } 63 | if v, p := a.(int16); p { 64 | return strconv.Itoa(int(v)) 65 | } 66 | if v, p := a.(int32); p { 67 | return strconv.Itoa(int(v)) 68 | } 69 | if v, p := a.(uint); p { 70 | return strconv.Itoa(int(v)) 71 | } 72 | if v, p := a.(uint64); p { 73 | return strconv.Itoa(int(v)) 74 | } 75 | if v, p := a.(float32); p { 76 | return strconv.FormatFloat(float64(v), 'f', -1, 32) 77 | } 78 | if v, p := a.(float64); p { 79 | return strconv.FormatFloat(v, 'f', -1, 32) 80 | } 81 | return "" 82 | } 83 | 84 | func Atoi(s string) uint64 { 85 | i, err := strconv.ParseUint(s, 10, 32) 86 | if err != nil { 87 | return 0 88 | } 89 | return i 90 | } 91 | 92 | func Atois(s []string) []uint64 { 93 | res := make([]uint64, len(s)) 94 | for i, str := range s { 95 | res[i] = Atoi(str) 96 | } 97 | return res 98 | } 99 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/ShugetsuSoft/pixivel-back 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/PuerkitoBio/goquery v1.7.1 // indirect 7 | github.com/RedisBloom/redisbloom-go v1.0.0 8 | github.com/ShugetsuSoft/loki-client-go v0.0.0-20220122062854-8cfb0cca9636 9 | github.com/antchfx/htmlquery v1.2.3 // indirect 10 | github.com/antchfx/xmlquery v1.3.6 // indirect 11 | github.com/gin-contrib/cors v1.3.1 12 | github.com/gin-gonic/gin v1.8.1 13 | github.com/go-playground/validator/v10 v10.11.1 // indirect 14 | github.com/gobwas/glob v0.2.3 // indirect 15 | github.com/goccy/go-json v0.9.11 // indirect 16 | github.com/gocolly/colly v1.2.0 17 | github.com/gomodule/redigo v1.8.5 18 | github.com/kennygrant/sanitize v1.2.4 // indirect 19 | github.com/klauspost/compress v1.11.13 // indirect 20 | github.com/mattn/go-isatty v0.0.16 // indirect 21 | github.com/mitchellh/mapstructure v1.1.2 22 | github.com/olivere/elastic/v7 v7.0.27 23 | github.com/pelletier/go-toml/v2 v2.0.5 // indirect 24 | github.com/prometheus/client_golang v1.12.0 25 | github.com/rabbitmq/amqp091-go v0.0.0-20210812094702-b2a427eb7d17 26 | github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca // indirect 27 | github.com/temoto/robotstxt v1.1.2 // indirect 28 | github.com/vmihailenco/msgpack/v5 v5.3.4 29 | go.mongodb.org/mongo-driver v1.7.1 30 | golang.org/x/crypto v0.1.0 // indirect 31 | google.golang.org/appengine v1.6.7 // indirect 32 | google.golang.org/grpc v1.33.2 33 | google.golang.org/protobuf v1.28.1 34 | gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect 35 | gopkg.in/yaml.v3 v3.0.1 36 | ) 37 | -------------------------------------------------------------------------------- /modules/responser/reader/dbreader.go: -------------------------------------------------------------------------------- 1 | package reader 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/ShugetsuSoft/pixivel-back/common/convert" 8 | "github.com/ShugetsuSoft/pixivel-back/common/models" 9 | ) 10 | 11 | func (r *Reader) IllustResponse(ctx context.Context, illustId uint64, forcefetch bool) (*models.IllustResponse, error) { 12 | retry := 1 13 | START: 14 | illust, err := r.dbops.QueryIllust(ctx, illustId, false) 15 | if err != nil { 16 | return nil, err 17 | } 18 | 19 | if illust == nil || forcefetch && time.Now().After(illust.UpdateTime.Add(time.Hour*24*2)) { 20 | if r.mode == models.ArchiveMode { 21 | return nil, models.ErrorArchiveMode 22 | } 23 | if retry == 0 { 24 | return nil, models.ErrorRetrivingFinishedTask 25 | } 26 | err = r.gen.IllustDetailTask(ctx, illustId) 27 | if err != nil { 28 | return nil, err 29 | } 30 | retry-- 31 | goto START 32 | 33 | } 34 | 35 | userId := uint64(illust.User) 36 | user, err := r.dbops.QueryUser(ctx, userId, false) 37 | if err != nil { 38 | return nil, err 39 | } 40 | 41 | if user == nil { 42 | if r.mode == models.ArchiveMode { 43 | return nil, models.ErrorArchiveMode 44 | } 45 | if retry == 0 { 46 | return nil, models.ErrorRetrivingFinishedTask 47 | } 48 | err = r.gen.UserDetailTask(ctx, userId) 49 | if err != nil { 50 | return nil, err 51 | } 52 | retry-- 53 | goto START 54 | } 55 | 56 | return convert.Illust2IllustResponse(illust, user), nil 57 | } 58 | 59 | func (r *Reader) UgoiraResponse(ctx context.Context, ugoiraId uint64, forcefetch bool) (*models.UgoiraResponse, error) { 60 | retry := 1 61 | START: 62 | ugoira, err := r.dbops.QueryUgoira(ctx, ugoiraId) 63 | if err != nil { 64 | return nil, err 65 | } 66 | 67 | if ugoira == nil || forcefetch && time.Now().After(ugoira.UpdateTime.Add(time.Hour*24*2)) { 68 | if r.mode == models.ArchiveMode { 69 | return nil, models.ErrorArchiveMode 70 | } 71 | if retry == 0 { 72 | return nil, models.ErrorRetrivingFinishedTask 73 | } 74 | err = r.gen.UgoiraDetailTask(ctx, ugoiraId) 75 | if err != nil { 76 | return nil, err 77 | } 78 | retry-- 79 | goto START 80 | } 81 | 82 | return convert.Ugoira2UgoiraResponse(ugoira), nil 83 | } 84 | 85 | func (r *Reader) IllustsResponse(ctx context.Context, illustIds []uint64) (*models.IllustsResponse, error) { 86 | var err error 87 | var illust *models.Illust 88 | illusts := make([]models.Illust, 0, len(illustIds)) 89 | for _, i := range illustIds { 90 | illust, err = r.dbops.QueryIllust(ctx, illustIds[i], false) 91 | if err != nil { 92 | return nil, err 93 | } 94 | if illust == nil { 95 | continue 96 | } 97 | illusts = append(illusts, *illust) 98 | } 99 | return convert.Illusts2IllustsResponse(illusts, false), nil 100 | } 101 | 102 | func (r *Reader) UserDetailResponse(ctx context.Context, userId uint64) (*models.UserResponse, error) { 103 | retry := 1 104 | START: 105 | user, err := r.dbops.QueryUser(ctx, userId, false) 106 | if err != nil { 107 | return nil, err 108 | } 109 | 110 | if user == nil { 111 | if r.mode == models.ArchiveMode { 112 | return nil, models.ErrorArchiveMode 113 | } 114 | if retry == 0 { 115 | return nil, models.ErrorRetrivingFinishedTask 116 | } 117 | err = r.gen.UserDetailTask(ctx, userId) 118 | if err != nil { 119 | return nil, err 120 | } 121 | retry-- 122 | goto START 123 | } 124 | 125 | return convert.User2UserResponse(user), nil 126 | } 127 | 128 | func (r *Reader) UserIllustsResponse(ctx context.Context, userId uint64, page int64, limit int64) (*models.IllustsResponse, error) { 129 | retry := 1 130 | START: 131 | user, err := r.dbops.QueryUser(ctx, userId, false) 132 | if err != nil { 133 | return nil, err 134 | } 135 | 136 | if user == nil { 137 | if r.mode == models.ArchiveMode { 138 | return nil, models.ErrorArchiveMode 139 | } 140 | if retry == 0 { 141 | return nil, models.ErrorRetrivingFinishedTask 142 | } 143 | err = r.gen.UserDetailTask(ctx, userId) 144 | if err != nil { 145 | return nil, err 146 | } 147 | err = r.gen.UserIllustsTask(ctx, userId) 148 | if err != nil { 149 | return nil, err 150 | } 151 | retry-- 152 | goto START 153 | } 154 | 155 | if (time.Now().After(user.IllustsUpdateTime.Add(time.Hour*24*4)) || user.IllustsCount == 0) && r.mode != models.ArchiveMode { 156 | if retry == 0 { 157 | goto END 158 | } 159 | err = r.gen.UserIllustsTask(ctx, userId) 160 | if err != nil { 161 | return nil, err 162 | } 163 | retry-- 164 | goto START 165 | } 166 | 167 | END: 168 | illusts, err := r.dbops.QueryIllustByUserWithPage(ctx, userId, page, limit, false) 169 | if err != nil { 170 | return nil, err 171 | } 172 | 173 | return convert.Illusts2IllustsResponse(illusts, user.IllustsCount > uint(limit*(page+1))), err 174 | } 175 | 176 | func (r *Reader) RankIllustsResponse(ctx context.Context, mode string, date string, page int, content string, limit int) (*models.IllustsResponse, error) { 177 | results, err := r.dbops.QueryRankIllusts(ctx, mode, date, content, page, limit) 178 | if err != nil { 179 | return nil, err 180 | } 181 | return convert.RankAggregateResult2IllustsResponses(results, page < 9 && len(results) != 0), nil 182 | } 183 | 184 | func (r *Reader) SampleIllustsResponse(ctx context.Context, quality int, limit int) (*models.IllustsResponse, error) { 185 | results, err := r.dbops.GetSampleIllusts(ctx, quality, limit, false) 186 | if err != nil { 187 | return nil, err 188 | } 189 | return convert.Illusts2IllustsResponse(results, false), nil 190 | } 191 | 192 | func (r *Reader) SampleUsersResponse(ctx context.Context, limit int) (*models.UsersResponse, error) { 193 | results, err := r.dbops.GetSampleUsers(ctx, limit, false) 194 | if err != nil { 195 | return nil, err 196 | } 197 | return convert.Users2UsersResponse(results, false), nil 198 | } 199 | -------------------------------------------------------------------------------- /modules/responser/reader/reader.go: -------------------------------------------------------------------------------- 1 | package reader 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 5 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 6 | "github.com/ShugetsuSoft/pixivel-back/common/models" 7 | "github.com/ShugetsuSoft/pixivel-back/modules/responser/task" 8 | ) 9 | 10 | type Reader struct { 11 | dbops *operations.DatabaseOperations 12 | gen *task.TaskGenerator 13 | mode models.Modes 14 | //shops *operations.SearchOperations 15 | } 16 | 17 | func NewReader(dbops *operations.DatabaseOperations, mq models.MessageQueue, taskchaname string, retrys uint, tracer *tasktracer.TaskTracer, mode models.Modes) *Reader { 18 | gen := task.NewTaskGenerator(mq, taskchaname, retrys, tracer) 19 | return &Reader{ 20 | dbops: dbops, 21 | gen: gen, 22 | mode: mode, 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /modules/responser/reader/searchreader.go: -------------------------------------------------------------------------------- 1 | package reader 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/ShugetsuSoft/pixivel-back/common/convert" 7 | "github.com/ShugetsuSoft/pixivel-back/common/models" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 9 | ) 10 | 11 | func (r *Reader) SearchIllustsResponse(ctx context.Context, keyword string, page int, limit int, sortpopularity bool, sortdate bool) (*models.IllustsSearchResponse, error) { 12 | illusts, hits, scores, highlights, err := r.dbops.SearchIllust(ctx, keyword, page, limit, sortpopularity, sortdate, false) 13 | if err != nil { 14 | return nil, err 15 | } 16 | 17 | return convert.Illusts2IllustsSearchResponse(illusts, hits > int64(limit*(page+1)), scores, highlights), err 18 | } 19 | 20 | func (r *Reader) SearchIllustsSuggestResponse(ctx context.Context, keyword string) (*models.SearchSuggestResponse, error) { 21 | suggests, err := r.dbops.SearchIllustSuggest(ctx, keyword) 22 | if err != nil { 23 | return nil, err 24 | } 25 | 26 | return &models.SearchSuggestResponse{ 27 | SuggestWords: suggests, 28 | }, nil 29 | } 30 | 31 | func (r *Reader) SearchUsersResponse(ctx context.Context, keyword string, page int, limit int) (*models.UsersSearchResponse, error) { 32 | users, hits, scores, highlights, err := r.dbops.SearchUser(ctx, keyword, page, limit, false) 33 | if err != nil { 34 | return nil, err 35 | } 36 | 37 | return convert.Users2UsersSearchResponse(users, hits > int64(limit*(page+1)), scores, highlights), err 38 | } 39 | 40 | func (r *Reader) SearchUsersSuggestResponse(ctx context.Context, keyword string) (*models.SearchSuggestResponse, error) { 41 | suggests, err := r.dbops.SearchUserSuggest(ctx, keyword) 42 | if err != nil { 43 | return nil, err 44 | } 45 | 46 | return &models.SearchSuggestResponse{ 47 | SuggestWords: suggests, 48 | }, nil 49 | } 50 | 51 | func (r *Reader) SearchTagsSuggestResponse(ctx context.Context, keyword string) (*models.SearchSuggestTagsResponse, error) { 52 | suggests, err := r.dbops.SearchTagSuggest(ctx, keyword) 53 | if err != nil { 54 | return nil, err 55 | } 56 | 57 | return &models.SearchSuggestTagsResponse{ 58 | SuggestTags: convert.Tags2TagResponses(suggests), 59 | }, nil 60 | } 61 | 62 | func (r *Reader) SearchIllustsByTagsResponse(ctx context.Context, tags []string, perfectmatch bool, page int, limit int, sortpopularity bool, sortdate bool) (*models.IllustsResponse, error) { 63 | if perfectmatch { 64 | illusts, err := r.dbops.QueryIllustsByTags(ctx, tags, page, limit, sortpopularity, sortdate, false) 65 | if err != nil { 66 | return nil, err 67 | } 68 | 69 | return convert.Illusts2IllustsResponse(illusts, len(illusts) >= limit), nil 70 | } 71 | return nil, nil 72 | } 73 | 74 | func (r *Reader) RecommendIllustsByIllustId(ctx context.Context, illustId uint64, k int) ([]models.Illust, error) { 75 | illusts, err := r.dbops.RecommendIllustsByIllustId(ctx, illustId, k, config.RecommendDrift, false) 76 | if err != nil { 77 | return nil, err 78 | } 79 | return illusts, nil 80 | } 81 | -------------------------------------------------------------------------------- /modules/responser/responser.go: -------------------------------------------------------------------------------- 1 | package responser 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/database/drivers" 5 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 6 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 7 | "github.com/ShugetsuSoft/pixivel-back/common/models" 8 | "github.com/gin-contrib/cors" 9 | "github.com/gin-gonic/gin" 10 | "time" 11 | ) 12 | 13 | type Responser struct { 14 | app *gin.Engine 15 | addr string 16 | } 17 | 18 | func NewResponser(addr string, dbops *operations.DatabaseOperations, mq models.MessageQueue, taskchaname string, retrys uint, tracer *tasktracer.TaskTracer, redis *drivers.RedisPool, debug bool, modes models.Modes, enableForceFetch bool) *Responser { 19 | app := gin.New() 20 | app.Use(cors.New(cors.Config{ 21 | AllowOrigins: []string{"https://pixivel.moe", "https://beta.pixivel.moe", "https://pixivel.art"}, 22 | AllowMethods: []string{"GET", "POST", "OPTIONS"}, 23 | AllowHeaders: []string{"Origin", "Content-Type", "User-Agent"}, 24 | ExposeHeaders: []string{"Content-Length"}, 25 | AllowCredentials: true, 26 | MaxAge: 12 * time.Hour, 27 | })) 28 | 29 | router := NewRouter(dbops, mq, taskchaname, retrys, tracer, redis, debug, modes, enableForceFetch) 30 | router.mount(app) 31 | 32 | return &Responser{ 33 | app, 34 | addr, 35 | } 36 | } 37 | 38 | func (res *Responser) Run() error { 39 | return res.app.Run(res.addr) 40 | } 41 | -------------------------------------------------------------------------------- /modules/responser/router.go: -------------------------------------------------------------------------------- 1 | package responser 2 | 3 | import ( 4 | "fmt" 5 | "strconv" 6 | "strings" 7 | "time" 8 | 9 | "github.com/ShugetsuSoft/pixivel-back/common/convert" 10 | "github.com/ShugetsuSoft/pixivel-back/common/database/drivers" 11 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 12 | "github.com/prometheus/client_golang/prometheus" 13 | 14 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 15 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 16 | "github.com/ShugetsuSoft/pixivel-back/common/models" 17 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 18 | "github.com/ShugetsuSoft/pixivel-back/modules/responser/reader" 19 | "github.com/gin-gonic/gin" 20 | ) 21 | 22 | type Response struct { 23 | Error bool `json:"error"` 24 | Message string `json:"message"` 25 | Data interface{} `json:"data"` 26 | } 27 | 28 | type Router struct { 29 | reader *reader.Reader 30 | cache *drivers.Cache 31 | debug bool 32 | enableForceFetch bool 33 | } 34 | 35 | func success(data interface{}) *Response { 36 | return &Response{Error: false, Message: "", Data: data} 37 | } 38 | 39 | func NewRouter(dbops *operations.DatabaseOperations, mq models.MessageQueue, taskchaname string, retrys uint, tracer *tasktracer.TaskTracer, redis *drivers.RedisPool, debugflag bool, mode models.Modes, enableForceFetch bool) *Router { 40 | return &Router{ 41 | reader: reader.NewReader(dbops, mq, taskchaname, retrys, tracer, mode), 42 | cache: drivers.NewCache(redis), 43 | debug: debugflag, 44 | enableForceFetch: enableForceFetch, 45 | } 46 | } 47 | 48 | func fail(err string) *Response { 49 | return &Response{Error: true, Message: err, Data: nil} 50 | } 51 | 52 | func (r *Router) Fail(c *gin.Context, code int, err error) { 53 | errResp := "" 54 | report := true 55 | if r.debug { 56 | errResp = fmt.Sprintf("%s", err) 57 | } else { 58 | errResp = func() string { 59 | switch err { 60 | case models.ErrorNoResult: 61 | return "未返回结果。请检查你所访问的图片链接是否正确。" 62 | case models.ErrorItemBanned: 63 | return "该图片由于违反我们的服务政策被我们禁止访问。此禁止与Pixiv无关。" 64 | case models.ErrorRetrivingFinishedTask: 65 | return "后台任务失败。这应该与您没有关系，如果重复出现，请告知我们。" 66 | case models.ErrorTimeOut: 67 | return "图片信息获取超时。这应该与您没有关系，如果重复出现，请告知我们。" 68 | case models.ErrorArchiveMode: 69 | report = false 70 | return "全站当前处于归档模式。您的访问受限制。" 71 | default: 72 | switch err.Error() { 73 | case "尚无权限浏览该作品": 74 | return "该图片可能曾经存在，但已经被删除。" 75 | case "抱歉，您当前所寻找的个用户已经离开了pixiv, 或者这ID不存在。": 76 | return "该用户可能曾经存在，但已经被删除" 77 | case "Error Visited": 78 | return "该图片可能不存在。" 79 | default: 80 | return "未知错误。" 81 | } 82 | } 83 | }() 84 | } 85 | 86 | if report { 87 | realIp := c.ClientIP() 88 | telemetry.Log(telemetry.Label{"pos": "ResponseError", "ip": realIp}, fmt.Sprintf("%s", err)) 89 | } 90 | c.JSON(code, &Response{Error: true, Message: errResp, Data: nil}) 91 | } 92 | 93 | func (r *Router) GetIllustHandler(c *gin.Context) { 94 | ctx := c.Request.Context() 95 | 96 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "illust"}).Inc() 97 | id := utils.Atoi(c.Param("id")) 98 | if id == 0 { 99 | return 100 | } 101 | 102 | forcefetch := false 103 | if r.enableForceFetch { 104 | if i, e := strconv.ParseBool(c.Query("forcefetch")); i && e == nil { 105 | forcefetch = true 106 | } 107 | } 108 | 109 | if !forcefetch { 110 | cached, err := r.cache.Get("illust", utils.Itoa(id)) 111 | if err != nil { 112 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 113 | } 114 | if cached != nil { 115 | c.JSON(200, success(cached)) 116 | return 117 | } 118 | } 119 | 120 | illust, err := r.reader.IllustResponse(ctx, id, forcefetch) 121 | 122 | if err != nil { 123 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "illust"}).Inc() 124 | r.Fail(c, 500, err) 125 | return 126 | } 127 | 128 | c.JSON(200, success(illust)) 129 | err = r.cache.Set("illust", illust, 60*60*12, utils.Itoa(id)) 130 | if err != nil { 131 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 132 | } 133 | } 134 | 135 | func (r *Router) GetUgoiraHandler(c *gin.Context) { 136 | ctx := c.Request.Context() 137 | 138 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "ugoira"}).Inc() 139 | id := utils.Atoi(c.Param("id")) 140 | if id == 0 { 141 | return 142 | } 143 | 144 | forcefetch := false 145 | if r.enableForceFetch { 146 | if i, e := strconv.ParseBool(c.Query("forcefetch")); i && e == nil { 147 | forcefetch = true 148 | } 149 | } 150 | 151 | if !forcefetch { 152 | cached, err := r.cache.Get("ugoira", utils.Itoa(id)) 153 | if err != nil { 154 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 155 | } 156 | if cached != nil { 157 | c.JSON(200, success(cached)) 158 | return 159 | } 160 | } 161 | 162 | ugoira, err := r.reader.UgoiraResponse(ctx, id, forcefetch) 163 | 164 | if err != nil { 165 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "ugoira"}).Inc() 166 | r.Fail(c, 500, err) 167 | return 168 | } 169 | 170 | c.JSON(200, success(ugoira)) 171 | err = r.cache.Set("ugoira", ugoira, 60*60*12, utils.Itoa(id)) 172 | if err != nil { 173 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 174 | } 175 | } 176 | 177 | func (r *Router) GetUserDetailHandler(c *gin.Context) { 178 | ctx := c.Request.Context() 179 | 180 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "user"}).Inc() 181 | id := utils.Atoi(c.Param("id")) 182 | if id == 0 { 183 | return 184 | } 185 | 186 | cached, err := r.cache.Get("user", utils.Itoa(id)) 187 | if err != nil { 188 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 189 | } 190 | if cached != nil { 191 | c.JSON(200, success(cached)) 192 | return 193 | } 194 | 195 | user, err := r.reader.UserDetailResponse(ctx, id) 196 | 197 | if err != nil { 198 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "user"}).Inc() 199 | r.Fail(c, 500, err) 200 | return 201 | } 202 | 203 | c.JSON(200, success(user)) 204 | err = r.cache.Set("user", user, 60*60*12, utils.Itoa(id)) 205 | if err != nil { 206 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 207 | } 208 | } 209 | 210 | func (r *Router) GetUserIllustsHandler(c *gin.Context) { 211 | ctx := c.Request.Context() 212 | 213 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "user-illust"}).Inc() 214 | id := utils.Atoi(c.Param("id")) 215 | if id == 0 { 216 | return 217 | } 218 | 219 | page := utils.Atoi(c.Query("page")) 220 | if page < 0 { 221 | page = 0 222 | } 223 | 224 | limit := utils.Atoi(c.Query("limit")) 225 | if limit > 40 || limit < 1 { 226 | limit = 30 227 | } 228 | 229 | cached, err := r.cache.Get("user-illust", utils.Itoa(id), utils.Itoa(page), utils.Itoa(limit)) 230 | if err != nil { 231 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 232 | } 233 | if cached != nil { 234 | c.JSON(200, success(cached)) 235 | return 236 | } 237 | 238 | illusts, err := r.reader.UserIllustsResponse(ctx, id, int64(page), int64(limit)) 239 | 240 | if err != nil { 241 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "user-illust"}).Inc() 242 | r.Fail(c, 500, err) 243 | return 244 | } 245 | 246 | c.JSON(200, success(illusts)) 247 | err = r.cache.Set("user-illust", illusts, 60*60*12, utils.Itoa(id), utils.Itoa(page), utils.Itoa(limit)) 248 | if err != nil { 249 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 250 | } 251 | } 252 | 253 | func (r *Router) SearchIllustHandler(c *gin.Context) { 254 | c.JSON(500, gin.H{ 255 | "error": "暂时不可用", 256 | }) 257 | return 258 | 259 | ctx := c.Request.Context() 260 | 261 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "search-illust"}).Inc() 262 | keyword := c.Param("keyword") 263 | if keyword == "" { 264 | return 265 | } 266 | 267 | page := utils.Atoi(c.Query("page")) 268 | if page < 0 { 269 | page = 0 270 | } 271 | 272 | limit := utils.Atoi(c.Query("limit")) 273 | if limit > 40 || limit < 1 { 274 | limit = 30 275 | } 276 | 277 | sortpop := false 278 | if i, e := strconv.ParseBool(c.Query("sortpop")); i && e == nil { 279 | sortpop = true 280 | } 281 | 282 | sortdate := false 283 | if i, e := strconv.ParseBool(c.Query("sortdate")); i && e == nil { 284 | sortdate = true 285 | } 286 | 287 | illusts, err := r.reader.SearchIllustsResponse(ctx, keyword, int(page), int(limit), sortpop, sortdate) 288 | 289 | if err != nil { 290 | if err == models.ErrorNoResult { 291 | r.Fail(c, 200, err) 292 | return 293 | } 294 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "search-illust"}).Inc() 295 | r.Fail(c, 500, err) 296 | return 297 | } 298 | 299 | c.JSON(200, success(illusts)) 300 | } 301 | 302 | func (r *Router) SearchIllustSuggestHandler(c *gin.Context) { 303 | c.JSON(500, gin.H{ 304 | "error": "暂时不可用", 305 | }) 306 | return 307 | 308 | ctx := c.Request.Context() 309 | 310 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "search-illust-suggest"}).Inc() 311 | keyword := c.Param("keyword") 312 | if keyword == "" { 313 | return 314 | } 315 | 316 | suggests, err := r.reader.SearchIllustsSuggestResponse(ctx, keyword) 317 | 318 | if err != nil { 319 | if err == models.ErrorNoResult { 320 | r.Fail(c, 200, err) 321 | return 322 | } 323 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "search-illust-suggest"}).Inc() 324 | r.Fail(c, 500, err) 325 | return 326 | } 327 | 328 | c.JSON(200, success(suggests)) 329 | } 330 | 331 | func (r *Router) SearchUserHandler(c *gin.Context) { 332 | c.JSON(500, gin.H{ 333 | "error": "暂时不可用", 334 | }) 335 | return 336 | 337 | ctx := c.Request.Context() 338 | 339 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "search-user"}).Inc() 340 | keyword := c.Param("keyword") 341 | if keyword == "" { 342 | return 343 | } 344 | 345 | page := utils.Atoi(c.Query("page")) 346 | if page < 0 { 347 | page = 0 348 | } 349 | 350 | limit := utils.Atoi(c.Query("limit")) 351 | if limit > 40 || limit < 1 { 352 | limit = 30 353 | } 354 | 355 | users, err := r.reader.SearchUsersResponse(ctx, keyword, int(page), int(limit)) 356 | 357 | if err != nil { 358 | if err == models.ErrorNoResult { 359 | r.Fail(c, 200, err) 360 | return 361 | } 362 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "search-user"}).Inc() 363 | r.Fail(c, 500, err) 364 | return 365 | } 366 | 367 | c.JSON(200, success(users)) 368 | } 369 | 370 | func (r *Router) SearchUserSuggestHandler(c *gin.Context) { 371 | c.JSON(500, gin.H{ 372 | "error": "暂时不可用", 373 | }) 374 | return 375 | 376 | ctx := c.Request.Context() 377 | 378 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "search-user-suggest"}).Inc() 379 | keyword := c.Param("keyword") 380 | if keyword == "" { 381 | return 382 | } 383 | 384 | suggests, err := r.reader.SearchUsersSuggestResponse(ctx, keyword) 385 | 386 | if err != nil { 387 | if err == models.ErrorNoResult { 388 | r.Fail(c, 200, err) 389 | return 390 | } 391 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "search-user-suggest"}).Inc() 392 | r.Fail(c, 500, err) 393 | return 394 | } 395 | 396 | c.JSON(200, success(suggests)) 397 | } 398 | 399 | func (r *Router) SearchTagSuggestHandler(c *gin.Context) { 400 | c.JSON(500, gin.H{ 401 | "error": "暂时不可用", 402 | }) 403 | return 404 | 405 | ctx := c.Request.Context() 406 | 407 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "search-tag-suggest"}).Inc() 408 | keyword := c.Param("keyword") 409 | if keyword == "" { 410 | return 411 | } 412 | 413 | suggests, err := r.reader.SearchTagsSuggestResponse(ctx, keyword) 414 | 415 | if err != nil { 416 | if err == models.ErrorNoResult { 417 | r.Fail(c, 200, err) 418 | return 419 | } 420 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "search-tag-suggest"}).Inc() 421 | r.Fail(c, 500, err) 422 | return 423 | } 424 | 425 | c.JSON(200, success(suggests)) 426 | } 427 | 428 | func (r *Router) SearchIllustByTagHandler(c *gin.Context) { 429 | c.JSON(500, gin.H{ 430 | "error": "暂时不可用", 431 | }) 432 | return 433 | 434 | ctx := c.Request.Context() 435 | 436 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "search-illust-by-tag"}).Inc() 437 | keywords := c.Param("keyword") 438 | if keywords == "" { 439 | return 440 | } 441 | tags := strings.Split(keywords, ",") 442 | 443 | page := utils.Atoi(c.Query("page")) 444 | if page < 0 { 445 | page = 0 446 | } 447 | 448 | limit := utils.Atoi(c.Query("limit")) 449 | if limit > 40 || limit < 1 { 450 | limit = 30 451 | } 452 | 453 | sortpop := false 454 | if i, e := strconv.ParseBool(c.Query("sortpop")); i && e == nil { 455 | sortpop = true 456 | } 457 | 458 | sortdate := false 459 | if i, e := strconv.ParseBool(c.Query("sortdate")); i && e == nil { 460 | sortdate = true 461 | } 462 | 463 | perfectmatch := true 464 | if i, e := strconv.ParseBool(c.Query("perfectmatch")); !i && e == nil { 465 | perfectmatch = false 466 | } 467 | 468 | illusts, err := r.reader.SearchIllustsByTagsResponse(ctx, tags, perfectmatch, int(page), int(limit), sortpop, sortdate) 469 | 470 | if err != nil { 471 | if err == models.ErrorNoResult { 472 | r.Fail(c, 200, err) 473 | return 474 | } 475 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "search-illust-by-tag"}).Inc() 476 | r.Fail(c, 500, err) 477 | return 478 | } 479 | 480 | c.JSON(200, success(illusts)) 481 | } 482 | 483 | func (r *Router) GetIllustsHandler(c *gin.Context) { 484 | ctx := c.Request.Context() 485 | 486 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "illusts"}).Inc() 487 | keywords := c.Param("ids") 488 | if keywords == "" { 489 | return 490 | } 491 | illustsstr := strings.Split(keywords, ",") 492 | if len(illustsstr) > 100 { 493 | c.JSON(400, fail("Query is too Large.")) 494 | return 495 | } 496 | illustsids := make([]uint64, len(illustsstr)) 497 | for i, id := range illustsstr { 498 | ida := utils.Atoi(id) 499 | if ida == 0 { 500 | return 501 | } 502 | illustsids[i] = ida 503 | } 504 | 505 | illusts, err := r.reader.IllustsResponse(ctx, illustsids) 506 | 507 | if err != nil { 508 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "illusts"}).Inc() 509 | r.Fail(c, 500, err) 510 | return 511 | } 512 | 513 | c.JSON(200, success(illusts)) 514 | } 515 | 516 | func (r *Router) RecommendIllustsByIllustIdHandler(c *gin.Context) { 517 | ctx := c.Request.Context() 518 | 519 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "illust-recommend"}).Inc() 520 | const maxpage = 5 521 | id := utils.Atoi(c.Param("id")) 522 | if id == 0 { 523 | return 524 | } 525 | 526 | page := utils.Atoi(c.Query("page")) 527 | if page < 0 { 528 | page = 0 529 | } 530 | 531 | if page >= maxpage { 532 | c.JSON(400, fail("没有更多了~")) 533 | return 534 | } 535 | 536 | cached, err := r.cache.Get("illust-recommend", utils.Itoa(id), utils.Itoa(page)) 537 | if err != nil { 538 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 539 | } 540 | if cached != nil { 541 | c.JSON(200, success(cached)) 542 | return 543 | } 544 | 545 | const limit = 30 546 | 547 | illusts, err := r.reader.RecommendIllustsByIllustId(ctx, id, limit*maxpage) 548 | 549 | if err != nil { 550 | if err == models.ErrorNoResult { 551 | r.Fail(c, 200, err) 552 | return 553 | } 554 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "illust-recommend"}).Inc() 555 | r.Fail(c, 500, err) 556 | return 557 | } 558 | 559 | for i := 0; i < maxpage; i++ { 560 | if len(illusts) < limit*(i+1) { 561 | c.JSON(400, fail("没有更多了~")) 562 | return 563 | } 564 | pagenow := illusts[limit*i : limit*(i+1)] 565 | if len(illusts) < limit*(i+2) { 566 | err = r.cache.Set("illust-recommend", convert.Illusts2IllustsResponse(pagenow, false), 60*60*2, utils.Itoa(id), utils.Itoa(i)) 567 | break 568 | } 569 | err = r.cache.Set("illust-recommend", convert.Illusts2IllustsResponse(pagenow, i < maxpage-1), 60*60*2, utils.Itoa(id), utils.Itoa(i)) 570 | if err != nil { 571 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 572 | } 573 | } 574 | 575 | c.JSON(200, success(convert.Illusts2IllustsResponse(illusts[limit*page:limit*(page+1)], page < maxpage-1))) 576 | } 577 | 578 | func (r *Router) GetRankHandler(c *gin.Context) { 579 | ctx := c.Request.Context() 580 | 581 | timeNow := time.Now() 582 | timeZ := time.Date(timeNow.Year(), timeNow.Month(), timeNow.Day(), 0, 0, 0, 0, timeNow.Location()) 583 | if timeNow.Before(timeZ.Add(2 * time.Hour)) { 584 | c.JSON(400, fail("排行榜目前暂无数据")) 585 | return 586 | } 587 | 588 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "rank"}).Inc() 589 | mode := c.Query("mode") 590 | modes := map[string]bool{"daily": true, "weekly": true, "monthly": true, "rookie": true, "original": true, "male": true, "female": true} 591 | if _, ok := modes[mode]; !ok { 592 | c.JSON(400, fail("Illegal mode")) 593 | return 594 | } 595 | 596 | content := c.Query("content") 597 | contents := map[string]bool{"all": true, "illust": true, "manga": true, "ugoira": true} 598 | if _, ok := contents[content]; !ok { 599 | c.JSON(400, fail("Illegal content")) 600 | return 601 | } 602 | 603 | date := c.Query("date") 604 | dateI, err := time.Parse("20060102", date) 605 | if err != nil { 606 | c.JSON(400, fail(fmt.Sprintf("Time parse Error. %s", err))) 607 | return 608 | } 609 | 610 | if dateI.After(time.Now().AddDate(0, 0, -1)) { 611 | c.JSON(400, fail("Rank Info DNE. 你是不是傻啊")) 612 | return 613 | } 614 | 615 | page := utils.Atoi(c.Query("page")) 616 | if page < 0 { 617 | page = 0 618 | } 619 | 620 | if page > 9 { 621 | c.JSON(400, fail("没有更多了~")) 622 | return 623 | } 624 | 625 | limit := 50 626 | 627 | illusts, err := r.reader.RankIllustsResponse(ctx, mode, date, int(page), content, limit) 628 | 629 | if err != nil { 630 | if err == models.ErrorNoResult { 631 | r.Fail(c, 200, err) 632 | return 633 | } 634 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "rank"}).Inc() 635 | r.Fail(c, 500, err) 636 | return 637 | } 638 | 639 | if page == 0 && illusts.HasNext == false { 640 | c.JSON(400, fail("排行榜后台暂无数据，请与维护者联系。")) 641 | return 642 | } 643 | 644 | c.JSON(200, success(illusts)) 645 | } 646 | 647 | func (r *Router) GetSampleIllustsHandler(c *gin.Context) { 648 | ctx := c.Request.Context() 649 | 650 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "sample-illusts"}).Inc() 651 | page := utils.Atoi(c.Query("p")) 652 | if page > 20 || page < 0 { 653 | page = 0 654 | } 655 | 656 | cached, err := r.cache.Get("illust-sample", utils.Itoa(page)) 657 | if err != nil { 658 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 659 | } 660 | if cached != nil { 661 | c.JSON(200, success(cached)) 662 | return 663 | } 664 | 665 | illusts, err := r.reader.SampleIllustsResponse(ctx, 15000, 50) 666 | 667 | if err != nil { 668 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "sample-illusts"}).Inc() 669 | r.Fail(c, 500, err) 670 | return 671 | } 672 | 673 | err = r.cache.Set("illust-sample", illusts, 60*60*6, utils.Itoa(page)) 674 | if err != nil { 675 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 676 | } 677 | 678 | c.JSON(200, success(illusts)) 679 | } 680 | 681 | func (r *Router) GetSampleUsersHandler(c *gin.Context) { 682 | ctx := c.Request.Context() 683 | 684 | telemetry.RequestsCount.With(prometheus.Labels{"handler": "sample-users"}).Inc() 685 | page := utils.Atoi(c.Query("p")) 686 | if page > 20 || page < 0 { 687 | page = 0 688 | } 689 | 690 | cached, err := r.cache.Get("user-sample", utils.Itoa(page)) 691 | if err != nil { 692 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 693 | } 694 | if cached != nil { 695 | c.JSON(200, success(cached)) 696 | return 697 | } 698 | 699 | users, err := r.reader.SampleUsersResponse(ctx, 50) 700 | 701 | if err != nil { 702 | telemetry.RequestsErrorCount.With(prometheus.Labels{"handler": "sample-users"}).Inc() 703 | r.Fail(c, 500, err) 704 | return 705 | } 706 | 707 | err = r.cache.Set("user-sample", users, 60*60*6, utils.Itoa(page)) 708 | if err != nil { 709 | telemetry.Log(telemetry.Label{"pos": "cache"}, err.Error()) 710 | } 711 | 712 | c.JSON(200, success(users)) 713 | } 714 | 715 | func (r *Router) mount(rout *gin.Engine) { 716 | r1 := rout.Group("/v2") 717 | // pixiv 718 | 719 | pixiv := r1.Group("/pixiv") 720 | { 721 | pixiv.GET("/illust/:id", r.GetIllustHandler) 722 | pixiv.GET("/user/:id", r.GetUserDetailHandler) 723 | pixiv.GET("/user/:id/illusts", r.GetUserIllustsHandler) 724 | pixiv.GET("/illust/search/:keyword", r.SearchIllustHandler) 725 | pixiv.GET("/illust/search/:keyword/suggest", r.SearchIllustSuggestHandler) 726 | pixiv.GET("/user/search/:keyword", r.SearchUserHandler) 727 | pixiv.GET("/user/search/:keyword/suggest", r.SearchUserSuggestHandler) 728 | pixiv.GET("/tag/search/:keyword", r.SearchIllustByTagHandler) 729 | pixiv.GET("/tag/search/:keyword/suggest", r.SearchTagSuggestHandler) 730 | pixiv.GET("/illust/:id/recommend", r.RecommendIllustsByIllustIdHandler) 731 | //pixiv.GET("/illusts/:ids", r.GetIllustsHandler) 732 | pixiv.GET("/rank/", r.GetRankHandler) 733 | pixiv.GET("/illusts/sample", r.GetSampleIllustsHandler) 734 | pixiv.GET("/user/sample", r.GetSampleUsersHandler) 735 | pixiv.GET("/ugoira/:id", r.GetUgoiraHandler) 736 | } 737 | } 738 | -------------------------------------------------------------------------------- /modules/responser/task/gens.go: -------------------------------------------------------------------------------- 1 | package task 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "time" 7 | 8 | "github.com/ShugetsuSoft/pixivel-back/common/models" 9 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 10 | ) 11 | 12 | type stringmap map[string]string 13 | 14 | func (gen *TaskGenerator) IllustDetailTask(ctx context.Context, id uint64) error { 15 | tid, existed, err := gen.tracer.NewTaskGroup(stringmap{"iid": utils.Itoa(id)}, 1) 16 | if err != nil { 17 | return err 18 | } 19 | if !existed { 20 | task := models.CrawlTask{ 21 | Group: tid, 22 | Type: models.CrawlIllustDetail, 23 | Params: map[string]string{"id": utils.Itoa(id)}, 24 | RetryCount: gen.retrys, 25 | } 26 | err = gen.SendTask(task) 27 | if err != nil { 28 | gen.tracer.RemoveTaskGroup(tid) 29 | return err 30 | } 31 | } 32 | err = gen.tracerlistener.WaitFor(ctx, tid, time.Minute) 33 | if err != nil { 34 | gen.tracer.RemoveTaskGroup(tid) 35 | return err 36 | } 37 | return nil 38 | } 39 | 40 | func (gen *TaskGenerator) UgoiraDetailTask(ctx context.Context, id uint64) error { 41 | tid, existed, err := gen.tracer.NewTaskGroup(stringmap{"ugid": utils.Itoa(id)}, 1) 42 | if err != nil { 43 | return err 44 | } 45 | fmt.Println(id) 46 | if !existed { 47 | task := models.CrawlTask{ 48 | Group: tid, 49 | Type: models.CrawlUgoiraDetail, 50 | Params: map[string]string{"id": utils.Itoa(id)}, 51 | RetryCount: gen.retrys, 52 | } 53 | err = gen.SendTask(task) 54 | if err != nil { 55 | gen.tracer.RemoveTaskGroup(tid) 56 | return err 57 | } 58 | } 59 | err = gen.tracerlistener.WaitFor(ctx, tid, time.Minute) 60 | if err != nil { 61 | gen.tracer.RemoveTaskGroup(tid) 62 | return err 63 | } 64 | return nil 65 | } 66 | 67 | func (gen *TaskGenerator) UserDetailTask(ctx context.Context, id uint64) error { 68 | tid, existed, err := gen.tracer.NewTaskGroup(stringmap{"uid": utils.Itoa(id)}, 1) 69 | if err != nil { 70 | return err 71 | } 72 | fmt.Println(id) 73 | if !existed { 74 | task := models.CrawlTask{ 75 | Group: tid, 76 | Type: models.CrawlUserDetail, 77 | Params: map[string]string{"id": utils.Itoa(id)}, 78 | RetryCount: gen.retrys, 79 | } 80 | err = gen.SendTask(task) 81 | if err != nil { 82 | gen.tracer.RemoveTaskGroup(tid) 83 | return err 84 | } 85 | } 86 | err = gen.tracerlistener.WaitFor(ctx, tid, time.Minute) 87 | if err != nil { 88 | gen.tracer.RemoveTaskGroup(tid) 89 | return err 90 | } 91 | return nil 92 | } 93 | 94 | func (gen *TaskGenerator) UserIllustsTask(ctx context.Context, id uint64) error { 95 | tid, existed, err := gen.tracer.NewTaskGroup(stringmap{"uiid": utils.Itoa(id)}, 1) 96 | if err != nil { 97 | return err 98 | } 99 | fmt.Println(id) 100 | if !existed { 101 | task := models.CrawlTask{ 102 | Group: tid, 103 | Type: models.CrawlUserIllusts, 104 | Params: map[string]string{"id": utils.Itoa(id)}, 105 | RetryCount: gen.retrys, 106 | } 107 | err = gen.SendTask(task) 108 | if err != nil { 109 | gen.tracer.RemoveTaskGroup(tid) 110 | return err 111 | } 112 | } 113 | err = gen.tracerlistener.WaitFor(ctx, tid, time.Minute) 114 | if err != nil { 115 | gen.tracer.RemoveTaskGroup(tid) 116 | return err 117 | } 118 | return nil 119 | } 120 | 121 | func (gen *TaskGenerator) RankInitTask(ctx context.Context, mode string, date string, content string) error { 122 | tid, existed, err := gen.tracer.NewTaskGroup(stringmap{"rtmd": mode, "rtdt": date}, 1) 123 | if err != nil { 124 | return err 125 | } 126 | if !existed { 127 | task := models.CrawlTask{ 128 | Group: tid, 129 | Type: models.CrawlRankIllusts, 130 | Params: map[string]string{"mode": mode, "page": "1", "date": date, "content": content}, 131 | RetryCount: gen.retrys, 132 | } 133 | err = gen.SendTask(task) 134 | if err != nil { 135 | gen.tracer.RemoveTaskGroup(tid) 136 | return err 137 | } 138 | } 139 | err = gen.tracerlistener.WaitFor(ctx, tid, 5*time.Minute) 140 | if err != nil { 141 | gen.tracer.RemoveTaskGroup(tid) 142 | return err 143 | } 144 | return nil 145 | } 146 | -------------------------------------------------------------------------------- /modules/responser/task/task.go: -------------------------------------------------------------------------------- 1 | package task 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 7 | "github.com/ShugetsuSoft/pixivel-back/common/models" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 9 | ) 10 | 11 | type TaskGenerator struct { 12 | mq models.MessageQueue 13 | taskchaname string 14 | retrys uint 15 | tracer *tasktracer.TaskTracer 16 | tracerlistener *tasktracer.TaskListener 17 | } 18 | 19 | func NewTaskGenerator(mq models.MessageQueue, taskchaname string, retrys uint, tracer *tasktracer.TaskTracer) *TaskGenerator { 20 | return &TaskGenerator{ 21 | mq, 22 | taskchaname, 23 | retrys, 24 | tracer, 25 | tracer.NewListener(context.Background()), 26 | } 27 | } 28 | 29 | func (gen *TaskGenerator) SendTask(task models.CrawlTask) error { 30 | fmt.Println(task) 31 | bin, err := utils.MsgPack(task) 32 | if err != nil { 33 | return err 34 | } 35 | return gen.mq.Publish(gen.taskchaname, models.MQMessage{Data: bin, Priority: 5}) 36 | } -------------------------------------------------------------------------------- /modules/spider/apis/pixiv.go: -------------------------------------------------------------------------------- 1 | package apis 2 | 3 | import "fmt" 4 | 5 | const ( 6 | Domain = "https://www.pixiv.net" 7 | 8 | IllustDetail = "/ajax/illust/%s" 9 | UserDetail = "/ajax/user/%s" 10 | UserIllusts = "/ajax/user/%s/profile/all" 11 | RankIllusts = "/ranking.php" 12 | UgoiraDetail = "/ajax/illust/%s/ugoira_meta" 13 | ) 14 | 15 | func IllustDetailG(id string) string { 16 | return Domain + fmt.Sprintf(IllustDetail, id) + "?lang=zh&full=1" 17 | } 18 | 19 | func UserDetailG(id string) string { 20 | return Domain + fmt.Sprintf(UserDetail, id) + "?lang=zh&full=1" 21 | } 22 | 23 | func UserIllustsG(id string) string { 24 | return Domain + fmt.Sprintf(UserIllusts, id) + "?lang=zh&full=1" 25 | } 26 | 27 | func UgoiraDetailG(id string) string { 28 | return Domain + fmt.Sprintf(UgoiraDetail, id) + "?lang=zh" 29 | } 30 | 31 | func RankIllustsG(mode string, page string, date string, content string) string { 32 | // mode = 'daily', 'weekly', 'monthly', 'rookie', 'male', 'female' 33 | return Domain + RankIllusts + fmt.Sprintf("?mode=%s&p=%s&date=%s&format=json&lang=zh&full=1&content=%s", mode, page, date, content) 34 | } 35 | -------------------------------------------------------------------------------- /modules/spider/pipeline/hook.go: -------------------------------------------------------------------------------- 1 | package pipeline 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 5 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/scheduler" 6 | "github.com/gocolly/colly" 7 | ) 8 | 9 | func (pipe *Pipeline) Hook(c *colly.Collector, taskq *scheduler.TaskQueue) { 10 | c.OnResponse(func(r *colly.Response) { 11 | telemetry.Log(telemetry.Label{"pos": "SpiderCrawlURL"}, r.Request.URL.String()) 12 | pipe.Response(r, taskq) 13 | }) 14 | c.OnError(func(r *colly.Response, err error) { 15 | telemetry.Log(telemetry.Label{"pos": "SpiderCrawlURL"}, r.Request.URL.String()) 16 | if err != nil { 17 | telemetry.Log(telemetry.Label{"pos": "SpiderCrawl"}, err.Error()) 18 | pipe.Error(r, taskq) 19 | } 20 | }) 21 | } 22 | -------------------------------------------------------------------------------- /modules/spider/pipeline/pipeline.go: -------------------------------------------------------------------------------- 1 | package pipeline 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/models" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 6 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 7 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/scheduler" 8 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/storage" 9 | "log" 10 | ) 11 | 12 | type Pipeline struct { 13 | exchangename string 14 | storer *storage.BetterInMemoryStorage 15 | } 16 | 17 | func NewPipeline(exchangename string, storer *storage.BetterInMemoryStorage) *Pipeline { 18 | return &Pipeline{ 19 | exchangename: exchangename, 20 | storer: storer, 21 | } 22 | } 23 | 24 | func (pipe *Pipeline) Send(tid string, typeid uint, data interface{}, priority uint8, taskq *scheduler.TaskQueue) { 25 | res := models.CrawlResponse{ 26 | Group: tid, 27 | Type: typeid, 28 | Response: data, 29 | } 30 | binres, err := utils.MsgPack(res) 31 | if err == nil { 32 | err = taskq.Publish(pipe.exchangename, models.MQMessage{ 33 | Data: binres, 34 | Priority: priority, 35 | }) 36 | if err != nil { 37 | telemetry.Log(telemetry.Label{"pos": "SpiderPipeline"}, err.Error()) 38 | log.Fatal(err) 39 | } 40 | } else { 41 | telemetry.Log(telemetry.Label{"pos": "SpiderPipeline"}, err.Error()) 42 | log.Fatal(err) 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /modules/spider/pipeline/response.go: -------------------------------------------------------------------------------- 1 | package pipeline 2 | 3 | import ( 4 | "encoding/json" 5 | "errors" 6 | "github.com/ShugetsuSoft/pixivel-back/common/convert" 7 | "github.com/ShugetsuSoft/pixivel-back/common/models" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 9 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 10 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/scheduler" 11 | "github.com/gocolly/colly" 12 | ) 13 | 14 | func (pipe *Pipeline) Response(r *colly.Response, taskq *scheduler.TaskQueue) { 15 | t := r.Ctx.GetAny("Task") 16 | ack := r.Ctx.GetAny("Ack") 17 | priority := r.Ctx.GetAny("Priority") 18 | switch task := t.(type) { 19 | case *models.CrawlTask: 20 | switch ackid := ack.(type) { 21 | case uint64: 22 | switch priority := priority.(type) { 23 | case uint8: 24 | resp, err := pipe.Deal(task.Type, r.Body, task) 25 | if resp != nil || err == nil { 26 | pipe.Send(task.Group, task.Type, resp, priority, taskq) 27 | taskq.Ack(ackid) 28 | } else { 29 | taskq.Reject(ackid) 30 | if task.RetryCount > 0 { 31 | task.RetryCount -= 1 32 | if err == models.InternalErrorLoginNeeded { 33 | task.Params["login"] = "1" 34 | } 35 | taskq.Resend(task, priority) 36 | } else { 37 | telemetry.SpiderErrorTaskCount.Inc() 38 | pipe.Send(task.Group, models.CrawlError, &models.CrawlErrorResponse{ 39 | TaskType: task.Type, 40 | Params: task.Params, 41 | Message: err.Error(), 42 | }, priority, taskq) 43 | } 44 | } 45 | switch urihash := r.Ctx.GetAny("Uri").(type) { 46 | case uint64: 47 | pipe.storer.ClearVisited(urihash) 48 | } 49 | } 50 | } 51 | } 52 | } 53 | 54 | func (pipe *Pipeline) Error(r *colly.Response, taskq *scheduler.TaskQueue) { 55 | telemetry.Log(telemetry.Label{"pos": "SipderError"}, utils.StringOut(r.Body)) 56 | t := r.Ctx.GetAny("Task") 57 | ack := r.Ctx.GetAny("Ack") 58 | switch task := t.(type) { 59 | case *models.CrawlTask: 60 | switch ackid := ack.(type) { 61 | case uint64: 62 | taskq.Reject(ackid) 63 | priority := r.Ctx.GetAny("Priority") 64 | switch priority := priority.(type) { 65 | case uint8: 66 | if task.RetryCount > 0 { 67 | task.RetryCount -= 1 68 | switch urihash := r.Ctx.GetAny("Uri").(type) { 69 | case uint64: 70 | pipe.storer.ClearVisited(urihash) 71 | } 72 | taskq.Resend(task, priority) 73 | } else { 74 | telemetry.SpiderErrorTaskCount.Inc() 75 | var raw models.ErrorRawResponse 76 | err := json.Unmarshal(r.Body, &raw) 77 | if err != nil { 78 | telemetry.Log(telemetry.Label{"pos": "SipderError"}, err.Error()) 79 | pipe.Send(task.Group, models.CrawlError, &models.CrawlErrorResponse{ 80 | TaskType: task.Type, 81 | Params: task.Params, 82 | Message: utils.StringOut(r.Body), 83 | }, priority, taskq) 84 | } else { 85 | pipe.Send(task.Group, models.CrawlError, &models.CrawlErrorResponse{ 86 | TaskType: task.Type, 87 | Params: task.Params, 88 | Message: raw.Message, 89 | }, priority, taskq) 90 | } 91 | } 92 | } 93 | } 94 | } 95 | } 96 | 97 | func (pipe *Pipeline) Deal(tasktype uint, rawb []byte, task *models.CrawlTask) (interface{}, error) { 98 | switch tasktype { 99 | case models.CrawlIllustDetail: 100 | var raw models.IllustRawResponse 101 | err := json.Unmarshal(rawb, &raw) 102 | if err != nil { 103 | return nil, err 104 | } 105 | if raw.Error { 106 | return nil, errors.New(raw.Message) 107 | } 108 | if raw.Body.Urls.Original == "" { 109 | return nil, models.InternalErrorLoginNeeded 110 | } 111 | return convert.IllustRaw2Illust(&raw.Body), nil 112 | case models.CrawlUserDetail: 113 | var raw models.UserRawResponse 114 | err := json.Unmarshal(rawb, &raw) 115 | if err != nil { 116 | return nil, err 117 | } 118 | if raw.Error { 119 | return nil, errors.New(raw.Message) 120 | } 121 | return convert.UserRaw2User(&raw.Body), err 122 | case models.CrawlUserIllusts: 123 | var raw models.UserIllustsRawResponse 124 | err := json.Unmarshal(rawb, &raw) 125 | if err != nil { 126 | return nil, err 127 | } 128 | if raw.Error { 129 | return nil, errors.New(raw.Message) 130 | } 131 | return convert.UserIllusts2UserIllustsResponse(utils.Atoi(task.Params["id"]), &raw.Body), nil 132 | case models.CrawlRankIllusts: 133 | var raw models.RankIllustsRawResponse 134 | err := json.Unmarshal(rawb, &raw) 135 | if err != nil { 136 | return nil, err 137 | } 138 | return convert.RankIllusts2RankIllustsResponse(&raw), nil 139 | case models.CrawlUgoiraDetail: 140 | var raw models.UgoiraRawResponse 141 | err := json.Unmarshal(rawb, &raw) 142 | if err != nil { 143 | return nil, err 144 | } 145 | if raw.Error { 146 | return nil, errors.New(raw.Message) 147 | } 148 | return convert.UgoiraRaw2Ugoira(&raw.Body, utils.Atoi(task.Params["id"])), nil 149 | } 150 | return nil, nil 151 | } 152 | -------------------------------------------------------------------------------- /modules/spider/scheduler/scheduler.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/models" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 6 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 7 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 8 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/apis" 9 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/storage" 10 | "github.com/gocolly/colly" 11 | "log" 12 | "net/http" 13 | ) 14 | 15 | type Scheduler struct { 16 | storer *storage.BetterInMemoryStorage 17 | cookie *http.Cookie 18 | exchangename string 19 | } 20 | 21 | func NewScheduler(cookie *http.Cookie, storer *storage.BetterInMemoryStorage, exchangename string) *Scheduler { 22 | return &Scheduler{ 23 | storer: storer, 24 | cookie: cookie, 25 | exchangename: exchangename, 26 | } 27 | } 28 | 29 | func (sch *Scheduler) Schedule(c *colly.Collector, taskq *TaskQueue) error { 30 | for { 31 | newTask, tag, priority, err := taskq.GetTask() 32 | if err != nil { 33 | return err 34 | } 35 | if newTask == nil { 36 | return models.ErrorChannelClosed 37 | } 38 | uri, needlogin := ConstructRequest(newTask) 39 | uhash := storage.GetUrlHash(uri) 40 | c.DisableCookies() 41 | ctx := colly.NewContext() 42 | ctx.Put("Task", newTask) 43 | ctx.Put("Priority", priority) 44 | ctx.Put("Ack", tag) 45 | ctx.Put("Uri", uhash) 46 | header := http.Header{ 47 | "User-Agent": []string{config.UserAgent}, 48 | } 49 | if needlogin { 50 | header.Add("Cookie", sch.cookie.String()) 51 | } 52 | if uri != "" { 53 | telemetry.SpiderTaskCount.Inc() 54 | err = c.Request("GET", uri, nil, ctx, header) 55 | if err != nil { 56 | taskq.Reject(tag) 57 | telemetry.Log(telemetry.Label{"pos": "SpiderScheduler"}, err.Error()) 58 | if err != colly.ErrAlreadyVisited { 59 | if newTask.RetryCount > 0 { 60 | newTask.RetryCount -= 1 61 | sch.storer.ClearVisited(uhash) 62 | taskq.Resend(newTask, priority) 63 | } 64 | } else { 65 | res := models.CrawlResponse{ 66 | Group: newTask.Group, 67 | Type: models.CrawlError, 68 | Response: &models.CrawlErrorResponse{ 69 | TaskType: newTask.Type, 70 | Params: newTask.Params, 71 | Message: "Error Visited", 72 | }, 73 | } 74 | binRes, err := utils.MsgPack(res) 75 | if err == nil { 76 | err = taskq.Publish(sch.exchangename, models.MQMessage{ 77 | Data: binRes, 78 | Priority: priority, 79 | }) 80 | if err != nil { 81 | telemetry.Log(telemetry.Label{"pos": "SpiderScheduler"}, err.Error()) 82 | log.Fatal(err) 83 | } 84 | } else { 85 | telemetry.Log(telemetry.Label{"pos": "SpiderScheduler"}, err.Error()) 86 | log.Fatal(err) 87 | } 88 | } 89 | } 90 | } 91 | } 92 | } 93 | 94 | func ConstructRequest(task *models.CrawlTask) (string, bool) { 95 | switch task.Type { 96 | case models.CrawlIllustDetail: 97 | isLogin := false 98 | if v, ok := task.Params["login"]; ok && v == "1" { 99 | isLogin = true 100 | } 101 | return apis.IllustDetailG(task.Params["id"]), isLogin 102 | case models.CrawlUserDetail: 103 | return apis.UserDetailG(task.Params["id"]), false 104 | case models.CrawlUserIllusts: 105 | return apis.UserIllustsG(task.Params["id"]), true 106 | case models.CrawlRankIllusts: 107 | return apis.RankIllustsG(task.Params["mode"], task.Params["page"], task.Params["date"], task.Params["content"]), false 108 | case models.CrawlUgoiraDetail: 109 | return apis.UgoiraDetailG(task.Params["id"]), true 110 | } 111 | return "", false 112 | } 113 | -------------------------------------------------------------------------------- /modules/spider/scheduler/task.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/models" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 6 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 7 | "log" 8 | ) 9 | 10 | type TaskQueue struct { 11 | queue <-chan models.MQMessage 12 | mq models.MessageQueue 13 | channame string 14 | } 15 | 16 | func NewTaskQueue(mq models.MessageQueue, channame string) *TaskQueue { 17 | taskchan, err := mq.Consume(channame) 18 | if err != nil { 19 | telemetry.Log(telemetry.Label{"pos": "SpiderSchedulerError"}, err.Error()) 20 | log.Fatal(err) 21 | } 22 | return &TaskQueue{ 23 | mq: mq, 24 | queue: taskchan, 25 | channame: channame, 26 | } 27 | } 28 | 29 | func (taskq *TaskQueue) GetTask() (*models.CrawlTask, uint64, uint8, error) { 30 | data, ok := <-taskq.queue 31 | if ok { 32 | var task models.CrawlTask 33 | err := utils.MsgUnpack(data.Data, &task) 34 | if err != nil { 35 | return nil, 0, 0, err 36 | } 37 | return &task, data.Tag, data.Priority, nil 38 | } 39 | return nil, 0, 0, nil 40 | } 41 | 42 | func (taskq *TaskQueue) Resend(task *models.CrawlTask, priority uint8) error { 43 | bin, err := utils.MsgPack(task) 44 | if err != nil { 45 | return err 46 | } 47 | return taskq.mq.Publish(taskq.channame, models.MQMessage{ 48 | Data: bin, 49 | Priority: priority, 50 | }) 51 | } 52 | 53 | func (taskq *TaskQueue) Publish(name string, message models.MQMessage) error { 54 | return taskq.mq.PublishToExchange(name, message) 55 | } 56 | 57 | func (taskq *TaskQueue) Ack(ackid uint64) error { 58 | return taskq.mq.Ack(ackid) 59 | } 60 | 61 | func (taskq *TaskQueue) Reject(ackid uint64) error { 62 | return taskq.mq.Reject(ackid) 63 | } 64 | -------------------------------------------------------------------------------- /modules/spider/spider.go: -------------------------------------------------------------------------------- 1 | package spider 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 5 | storage2 "github.com/ShugetsuSoft/pixivel-back/modules/spider/storage" 6 | "log" 7 | "net/http" 8 | "time" 9 | 10 | "github.com/ShugetsuSoft/pixivel-back/common/models" 11 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/pipeline" 12 | "github.com/ShugetsuSoft/pixivel-back/modules/spider/scheduler" 13 | "github.com/gocolly/colly" 14 | ) 15 | 16 | type Spider struct { 17 | col *colly.Collector 18 | sche *scheduler.Scheduler 19 | pipe *pipeline.Pipeline 20 | mq models.MessageQueue 21 | inqueue string 22 | } 23 | 24 | func NewSpider(mq models.MessageQueue, inqueue string, outexchange string, loginss string, threads int) (*Spider, error) { 25 | col := colly.NewCollector( 26 | colly.Async(true), 27 | ) 28 | col.Limit(&colly.LimitRule{ 29 | DomainGlob: "*", 30 | Parallelism: threads, 31 | RandomDelay: 600 * time.Millisecond, 32 | }) 33 | 34 | storage := &storage2.BetterInMemoryStorage{} 35 | if err := col.SetStorage(storage); err != nil { 36 | return nil, err 37 | } 38 | 39 | cookie := &http.Cookie{ 40 | Name: "PHPSESSID", 41 | Value: loginss, 42 | } 43 | 44 | pipe := pipeline.NewPipeline(outexchange, storage) 45 | return &Spider{ 46 | mq: mq, 47 | inqueue: inqueue, 48 | col: col, 49 | sche: scheduler.NewScheduler(cookie, storage, outexchange), 50 | pipe: pipe, 51 | }, nil 52 | } 53 | 54 | func (s *Spider) Crawl() { 55 | BEGIN: 56 | taskq := scheduler.NewTaskQueue(s.mq, s.inqueue) 57 | s.pipe.Hook(s.col, taskq) 58 | err := s.sche.Schedule(s.col, taskq) 59 | if err != nil { 60 | if err == models.ErrorChannelClosed { 61 | goto BEGIN 62 | } 63 | telemetry.Log(telemetry.Label{"pos": "Spider"}, err.Error()) 64 | log.Fatal(err) 65 | } 66 | } 67 | -------------------------------------------------------------------------------- /modules/spider/storage/storage.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import ( 4 | "net/http" 5 | "net/http/cookiejar" 6 | "net/url" 7 | "strings" 8 | "sync" 9 | ) 10 | 11 | type BetterInMemoryStorage struct { 12 | visitedURLs map[uint64]bool 13 | lock *sync.RWMutex 14 | jar *cookiejar.Jar 15 | } 16 | 17 | // Init initializes InMemoryStorage 18 | func (s *BetterInMemoryStorage) Init() error { 19 | if s.visitedURLs == nil { 20 | s.visitedURLs = make(map[uint64]bool) 21 | } 22 | if s.lock == nil { 23 | s.lock = &sync.RWMutex{} 24 | } 25 | if s.jar == nil { 26 | var err error 27 | s.jar, err = cookiejar.New(nil) 28 | return err 29 | } 30 | return nil 31 | } 32 | 33 | // Visited implements Storage.Visited() 34 | func (s *BetterInMemoryStorage) Visited(requestID uint64) error { 35 | s.lock.Lock() 36 | s.visitedURLs[requestID] = true 37 | s.lock.Unlock() 38 | return nil 39 | } 40 | 41 | // IsVisited implements Storage.IsVisited() 42 | func (s *BetterInMemoryStorage) IsVisited(requestID uint64) (bool, error) { 43 | s.lock.RLock() 44 | visited := s.visitedURLs[requestID] 45 | s.lock.RUnlock() 46 | return visited, nil 47 | } 48 | 49 | func (s *BetterInMemoryStorage) ClearVisited(requestID uint64) error { 50 | s.lock.Lock() 51 | delete(s.visitedURLs, requestID) 52 | s.lock.Unlock() 53 | return nil 54 | } 55 | 56 | // Cookies implements Storage.Cookies() 57 | func (s *BetterInMemoryStorage) Cookies(u *url.URL) string { 58 | return StringifyCookies(s.jar.Cookies(u)) 59 | } 60 | 61 | // SetCookies implements Storage.SetCookies() 62 | func (s *BetterInMemoryStorage) SetCookies(u *url.URL, cookies string) { 63 | s.jar.SetCookies(u, UnstringifyCookies(cookies)) 64 | } 65 | 66 | // Close implements Storage.Close() 67 | func (s *BetterInMemoryStorage) Close() error { 68 | return nil 69 | } 70 | 71 | // StringifyCookies serializes list of http.Cookies to string 72 | func StringifyCookies(cookies []*http.Cookie) string { 73 | // Stringify cookies. 74 | cs := make([]string, len(cookies)) 75 | for i, c := range cookies { 76 | cs[i] = c.String() 77 | } 78 | return strings.Join(cs, "\n") 79 | } 80 | 81 | // UnstringifyCookies deserializes a cookie string to http.Cookies 82 | func UnstringifyCookies(s string) []*http.Cookie { 83 | h := http.Header{} 84 | for _, c := range strings.Split(s, "\n") { 85 | h.Add("Set-Cookie", c) 86 | } 87 | r := http.Response{Header: h} 88 | return r.Cookies() 89 | } 90 | 91 | // ContainsCookie checks if a cookie name is represented in cookies 92 | func ContainsCookie(cookies []*http.Cookie, name string) bool { 93 | for _, c := range cookies { 94 | if c.Name == name { 95 | return true 96 | } 97 | } 98 | return false 99 | } -------------------------------------------------------------------------------- /modules/spider/storage/utils.go: -------------------------------------------------------------------------------- 1 | package storage 2 | 3 | import "hash/fnv" 4 | 5 | func GetUrlHash(u string) uint64 { 6 | h := fnv.New64a() 7 | h.Write([]byte(u)) 8 | return h.Sum64() 9 | } -------------------------------------------------------------------------------- /modules/storer/handler.go: -------------------------------------------------------------------------------- 1 | package storer 2 | 3 | import ( 4 | "context" 5 | "github.com/ShugetsuSoft/pixivel-back/common/models" 6 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 7 | "github.com/ShugetsuSoft/pixivel-back/modules/storer/source" 8 | "github.com/mitchellh/mapstructure" 9 | ) 10 | 11 | func (st *Storer) handleDatabase(dataq *source.DataQueue) error { 12 | ctx := context.Background() 13 | 14 | for { 15 | data, tag, priority, err := dataq.GetData() 16 | if err != nil { 17 | return err 18 | } 19 | if data == nil { 20 | return models.ErrorChannelClosed 21 | } 22 | switch data.Type { 23 | case models.CrawlIllustDetail: 24 | var resdata models.Illust 25 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 26 | st.tracer.FailTask(data.Group, err.Error()) 27 | return err 28 | } 29 | 30 | err = st.ops.InsertIllust(ctx, &resdata) 31 | if err != nil { 32 | st.tracer.FailTask(data.Group, err.Error()) 33 | return err 34 | } 35 | 36 | st.tracer.FinishTask(data.Group) 37 | err = dataq.Ack(tag) 38 | if err != nil { 39 | return err 40 | } 41 | 42 | case models.CrawlUserDetail: 43 | var resdata models.User 44 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 45 | st.tracer.FailTask(data.Group, err.Error()) 46 | return err 47 | } 48 | 49 | err = st.ops.InsertUser(ctx, &resdata) 50 | if err != nil { 51 | st.tracer.FailTask(data.Group, err.Error()) 52 | return err 53 | } 54 | 55 | st.tracer.FinishTask(data.Group) 56 | err = dataq.Ack(tag) 57 | if err != nil { 58 | return err 59 | } 60 | case models.CrawlUserIllusts: 61 | var resdata models.UserIllustsResponse 62 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 63 | st.tracer.FailTask(data.Group, err.Error()) 64 | return err 65 | } 66 | illustCount := uint(len(resdata.Illusts)) 67 | 68 | user, err := st.ops.QueryUser(ctx, resdata.UserID, true) 69 | if err != nil { 70 | st.tracer.FailTask(data.Group, err.Error()) 71 | return err 72 | } 73 | 74 | if user != nil { 75 | if illustCount != user.IllustsCount { 76 | err = st.ops.SetIllustsCount(ctx, resdata.UserID, illustCount) 77 | if err != nil { 78 | st.tracer.FailTask(data.Group, err.Error()) 79 | return err 80 | } 81 | } 82 | } else { 83 | task := models.CrawlTask{ 84 | Group: "", 85 | Type: models.CrawlUserDetail, 86 | Params: map[string]string{"id": utils.Itoa(resdata.UserID)}, 87 | RetryCount: st.retrys, 88 | } 89 | err = st.task.SendTask(task, priority) 90 | if err != nil { 91 | st.tracer.FailTask(data.Group, err.Error()) 92 | return err 93 | } 94 | } 95 | 96 | err = st.ops.UpdateUserIllustsTime(ctx, resdata.UserID) 97 | if err != nil { 98 | st.tracer.FailTask(data.Group, err.Error()) 99 | return err 100 | } 101 | 102 | if user != nil && illustCount != user.IllustsCount || user == nil { 103 | for id := range resdata.Illusts { 104 | exist, err := st.ops.IsIllustExist(resdata.Illusts[id]) 105 | if err != nil { 106 | st.tracer.FailTask(data.Group, err.Error()) 107 | return err 108 | } 109 | if !exist { 110 | task := models.CrawlTask{ 111 | Group: data.Group, 112 | Type: models.CrawlIllustDetail, 113 | Params: map[string]string{"id": utils.Itoa(resdata.Illusts[id])}, 114 | RetryCount: st.retrys, 115 | } 116 | err = st.tracer.NewTask(data.Group) 117 | if err != nil { 118 | st.tracer.FailTask(data.Group, err.Error()) 119 | return err 120 | } 121 | err = st.task.SendTask(task, priority) 122 | if err != nil { 123 | st.tracer.FailTask(data.Group, err.Error()) 124 | return err 125 | } 126 | } 127 | } 128 | } 129 | 130 | st.tracer.FinishTask(data.Group) 131 | err = dataq.Ack(tag) 132 | if err != nil { 133 | return err 134 | } 135 | 136 | case models.CrawlRankIllusts: 137 | var resdata models.RankIllustsResponseMessage 138 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 139 | st.tracer.FailTask(data.Group, err.Error()) 140 | return err 141 | } 142 | if resdata.Page == 1 { 143 | _, err = st.ops.InsertRank(ctx, resdata.Mode, resdata.Date, resdata.Content) 144 | if err != nil { 145 | return err 146 | } 147 | } 148 | rankIllusts := make([]models.RankIllust, len(resdata.Illusts)) 149 | for i := range resdata.Illusts { 150 | rankIllusts[i].Rank = resdata.Illusts[i].Pos 151 | rankIllusts[i].ID = resdata.Illusts[i].ID 152 | task := models.CrawlTask{ 153 | Group: data.Group, 154 | Type: models.CrawlIllustDetail, 155 | Params: map[string]string{"id": utils.Itoa(resdata.Illusts[i].ID)}, 156 | RetryCount: st.retrys, 157 | } 158 | err = st.tracer.NewTask(data.Group) 159 | if err != nil { 160 | st.tracer.FailTask(data.Group, err.Error()) 161 | return err 162 | } 163 | err = st.task.SendTask(task, priority) 164 | if err != nil { 165 | st.tracer.FailTask(data.Group, err.Error()) 166 | return err 167 | } 168 | } 169 | err := st.ops.AddRankIllusts(ctx, resdata.Mode, resdata.Date, resdata.Content, rankIllusts) 170 | if err != nil { 171 | st.tracer.FailTask(data.Group, err.Error()) 172 | return err 173 | } 174 | if resdata.Next { 175 | task := models.CrawlTask{ 176 | Group: data.Group, 177 | Type: models.CrawlRankIllusts, 178 | Params: map[string]string{"mode": resdata.Mode, "page": utils.Itoa(resdata.Page + 1), "date": resdata.Date, "content": resdata.Content}, 179 | RetryCount: st.retrys, 180 | } 181 | err = st.tracer.NewTask(data.Group) 182 | if err != nil { 183 | st.tracer.FailTask(data.Group, err.Error()) 184 | return err 185 | } 186 | err = st.task.SendTask(task, priority) 187 | if err != nil { 188 | st.tracer.FailTask(data.Group, err.Error()) 189 | return err 190 | } 191 | } 192 | 193 | st.tracer.FinishTask(data.Group) 194 | err = dataq.Ack(tag) 195 | if err != nil { 196 | return err 197 | } 198 | case models.CrawlUgoiraDetail: 199 | var resdata models.Ugoira 200 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 201 | st.tracer.FailTask(data.Group, err.Error()) 202 | return err 203 | } 204 | 205 | err = st.ops.InsertUgoira(ctx, &resdata) 206 | if err != nil { 207 | st.tracer.FailTask(data.Group, err.Error()) 208 | return err 209 | } 210 | 211 | st.tracer.FinishTask(data.Group) 212 | err = dataq.Ack(tag) 213 | if err != nil { 214 | return err 215 | } 216 | 217 | case models.CrawlError: 218 | var errData models.CrawlErrorResponse 219 | if err := mapstructure.Decode(data.Response, &errData); err != nil { 220 | st.tracer.FailTask(data.Group, err.Error()) 221 | return err 222 | } 223 | switch errData.Message { 224 | case "抱歉，您当前所寻找的个用户已经离开了pixiv, 或者这ID不存在。": 225 | id := utils.Atoi(errData.Params["id"]) 226 | err = st.ops.DeleteUser(ctx, id) 227 | if err != nil { 228 | st.tracer.FailTask(data.Group, err.Error()) 229 | return err 230 | } 231 | case "尚无权限浏览该作品": 232 | id := utils.Atoi(errData.Params["id"]) 233 | err = st.ops.DeleteIllust(ctx, id) 234 | if err != nil { 235 | st.tracer.FailTask(data.Group, err.Error()) 236 | return err 237 | } 238 | } 239 | 240 | st.tracer.FailTask(data.Group, errData.Message) 241 | err = dataq.Ack(tag) 242 | if err != nil { 243 | return err 244 | } 245 | default: 246 | err = dataq.Ack(tag) 247 | if err != nil { 248 | return err 249 | } 250 | } 251 | } 252 | } 253 | 254 | func (st *Storer) handleElasticSearch(dataq *source.DataQueue) error { 255 | ctx := context.Background() 256 | 257 | for { 258 | data, tag, _, err := dataq.GetData() 259 | if err != nil { 260 | return err 261 | } 262 | if data == nil { 263 | return models.ErrorChannelClosed 264 | } 265 | switch data.Type { 266 | case models.CrawlIllustDetail: 267 | var resdata models.Illust 268 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 269 | return err 270 | } 271 | 272 | err = st.ops.InsertIllustSearch(ctx, &resdata) 273 | if err != nil { 274 | return err 275 | } 276 | 277 | err = dataq.Ack(tag) 278 | if err != nil { 279 | return err 280 | } 281 | 282 | case models.CrawlUserDetail: 283 | var resdata models.User 284 | if err := mapstructure.Decode(data.Response, &resdata); err != nil { 285 | return err 286 | } 287 | 288 | err = st.ops.InsertUserSearch(ctx, &resdata) 289 | if err != nil { 290 | return err 291 | } 292 | 293 | err = dataq.Ack(tag) 294 | if err != nil { 295 | return err 296 | } 297 | default: 298 | err = dataq.Ack(tag) 299 | if err != nil { 300 | return err 301 | } 302 | } 303 | } 304 | } 305 | -------------------------------------------------------------------------------- /modules/storer/source/queue.go: -------------------------------------------------------------------------------- 1 | package source 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/models" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 6 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 7 | "log" 8 | ) 9 | 10 | type DataQueue struct { 11 | queue <-chan models.MQMessage 12 | mq models.MessageQueue 13 | } 14 | 15 | func NewDataQueue(mq models.MessageQueue, channame string) *DataQueue { 16 | taskchan, err := mq.Consume(channame) 17 | if err != nil { 18 | telemetry.Log(telemetry.Label{"pos": "StorerQueue"}, err.Error()) 19 | log.Fatal(err) 20 | } 21 | return &DataQueue{ 22 | mq: mq, 23 | queue: taskchan, 24 | } 25 | } 26 | 27 | func (dataq *DataQueue) GetData() (*models.CrawlResponse, uint64, uint8, error) { 28 | data, ok := <-dataq.queue 29 | if ok { 30 | var task models.CrawlResponse 31 | err := utils.MsgUnpack(data.Data, &task) 32 | if err != nil { 33 | return nil, 0, 0, err 34 | } 35 | return &task, data.Tag, data.Priority, nil 36 | } 37 | return nil, 0, 0, nil 38 | } 39 | 40 | func (dataq *DataQueue) Publish(name string, message models.MQMessage) error { 41 | return dataq.mq.Publish(name, message) 42 | } 43 | 44 | func (dataq *DataQueue) Ack(ackid uint64) error { 45 | return dataq.mq.Ack(ackid) 46 | } 47 | 48 | func (dataq *DataQueue) Reject(ackid uint64) error { 49 | return dataq.mq.Reject(ackid) 50 | } 51 | -------------------------------------------------------------------------------- /modules/storer/source/task.go: -------------------------------------------------------------------------------- 1 | package source 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/models" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils" 6 | ) 7 | 8 | type Task struct { 9 | mq models.MessageQueue 10 | taskchaname string 11 | } 12 | 13 | func NewTaskGenerator(mq models.MessageQueue, taskchaname string) *Task { 14 | return &Task{ 15 | mq, 16 | taskchaname, 17 | } 18 | } 19 | 20 | func (gen *Task) SendTask(task models.CrawlTask, priority uint8) error { 21 | bin, err := utils.MsgPack(task) 22 | if err != nil { 23 | return err 24 | } 25 | return gen.mq.Publish(gen.taskchaname, models.MQMessage{Data: bin, Priority: priority}) 26 | } 27 | -------------------------------------------------------------------------------- /modules/storer/storer.go: -------------------------------------------------------------------------------- 1 | package storer 2 | 3 | import ( 4 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 5 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 6 | "github.com/ShugetsuSoft/pixivel-back/common/models" 7 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 8 | "github.com/ShugetsuSoft/pixivel-back/modules/storer/source" 9 | "log" 10 | ) 11 | 12 | type Storer struct { 13 | mq models.MessageQueue 14 | ops *operations.DatabaseOperations 15 | tracer *tasktracer.TaskTracer 16 | task *source.Task 17 | retrys uint 18 | } 19 | 20 | func NewStorer(mq models.MessageQueue, taskchanname string, retrys uint, ops *operations.DatabaseOperations, tracer *tasktracer.TaskTracer) *Storer { 21 | return &Storer{ 22 | mq: mq, 23 | ops: ops, 24 | tracer: tracer, 25 | task: source.NewTaskGenerator(mq, taskchanname), 26 | retrys: retrys, 27 | } 28 | } 29 | 30 | func (st *Storer) StoreDB(dbchanname string) { 31 | BEGIN: 32 | dataq := source.NewDataQueue(st.mq, dbchanname) 33 | var err error 34 | err = st.handleDatabase(dataq) 35 | if err != nil { 36 | if err == models.ErrorChannelClosed { 37 | goto BEGIN 38 | } 39 | telemetry.Log(telemetry.Label{"pos": "Storer"}, err.Error()) 40 | log.Fatal(err) 41 | } 42 | } 43 | 44 | func (st *Storer) StoreES(eschanname string) { 45 | BEGIN: 46 | dataq := source.NewDataQueue(st.mq, eschanname) 47 | var err error 48 | err = st.handleElasticSearch(dataq) 49 | if err != nil { 50 | if err == models.ErrorChannelClosed { 51 | goto BEGIN 52 | } 53 | telemetry.Log(telemetry.Label{"pos": "Storer"}, err.Error()) 54 | log.Fatal(err) 55 | } 56 | } 57 | -------------------------------------------------------------------------------- /modules/util.go: -------------------------------------------------------------------------------- 1 | package modules 2 | 3 | import ( 4 | "context" 5 | "io/ioutil" 6 | "os" 7 | 8 | "github.com/ShugetsuSoft/pixivel-back/common/models" 9 | "gopkg.in/yaml.v3" 10 | 11 | "github.com/ShugetsuSoft/pixivel-back/common/database/drivers" 12 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 13 | ) 14 | 15 | func ReadConfig(path string) (*models.Config, error) { 16 | var config models.Config 17 | confile, err := os.Open(path) 18 | if err != nil { 19 | return nil, err 20 | } 21 | defer confile.Close() 22 | content, err := ioutil.ReadAll(confile) 23 | if err != nil { 24 | return nil, err 25 | } 26 | err = yaml.Unmarshal(content, &config) 27 | return &config, err 28 | } 29 | 30 | func NewDB(conf *models.Config, ctx context.Context) (*drivers.MongoDatabase, error) { 31 | db, err := drivers.NewMongoDatabase(ctx, conf.Mongodb.URI, config.DatabaseName) 32 | if err != nil { 33 | return nil, err 34 | } 35 | return db, nil 36 | } 37 | 38 | func NewMQ(conf *models.Config) (*drivers.RabbitMQ, error) { 39 | mq, err := drivers.NewRabbitMQ(conf.Rabbitmq.URI) 40 | if err != nil { 41 | return nil, err 42 | } 43 | mq.QueueDeclare(config.CrawlTaskQueue) 44 | mq.QueueDeclare(config.CrawlResponsesMongodb) 45 | mq.QueueDeclare(config.CrawlResponsesElastic) 46 | mq.ExchangeDeclare(config.CrawlResponsesExchange, "fanout") 47 | mq.QueueBindExchange(config.CrawlResponsesMongodb, config.CrawlResponsesMongodb, config.CrawlResponsesExchange) 48 | mq.QueueBindExchange(config.CrawlResponsesElastic, config.CrawlResponsesElastic, config.CrawlResponsesExchange) 49 | return mq, nil 50 | } 51 | 52 | func NewES(conf *models.Config, ctx context.Context) (*drivers.ElasticSearch, error) { 53 | es, err := drivers.NewElasticSearchClient(ctx, conf.Elasticsearch.URI, conf.Elasticsearch.User, conf.Elasticsearch.Pass) 54 | return es, err 55 | } 56 | 57 | func NewRD(redisuri string) *drivers.RedisPool { 58 | redis := drivers.NewRedisPool(redisuri) 59 | return redis 60 | } 61 | 62 | func NewNB(conf *models.Config, ctx context.Context) (*drivers.NearDB, error) { 63 | ndb, err := drivers.NewNearDB(conf.Neardb.URI) 64 | return ndb, err 65 | } 66 | -------------------------------------------------------------------------------- /responser.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 7 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 9 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 10 | "github.com/ShugetsuSoft/pixivel-back/modules" 11 | "github.com/ShugetsuSoft/pixivel-back/modules/responser" 12 | "log" 13 | ) 14 | 15 | func main() { 16 | ctx, cancel := context.WithCancel(context.Background()) 17 | defer cancel() 18 | 19 | var configPath string 20 | flag.StringVar(&configPath, "config", "config.yaml", "Config File Path") 21 | flag.Parse() 22 | 23 | conf, err := modules.ReadConfig(configPath) 24 | if err != nil { 25 | log.Fatal(err) 26 | } 27 | 28 | telemetry.RegisterResponser() 29 | errchan := telemetry.RunLoki(conf.General.Loki, "responser") 30 | go func() { 31 | for err := range errchan { 32 | log.Println(err) 33 | } 34 | }() 35 | go telemetry.RunPrometheus(conf.General.Prometheus) 36 | 37 | db, err := modules.NewDB(conf, ctx) 38 | if err != nil { 39 | log.Fatal(err) 40 | } 41 | mq, err := modules.NewMQ(conf) 42 | if err != nil { 43 | log.Fatal(err) 44 | } 45 | ndb, err := modules.NewNB(conf, ctx) 46 | if err != nil { 47 | log.Fatal(err) 48 | } 49 | es, err := modules.NewES(conf, ctx) 50 | if err != nil { 51 | log.Fatal(err) 52 | } 53 | 54 | cacheRedis := modules.NewRD(conf.Redis.CacheRedis.URI) 55 | messageRedis := modules.NewRD(conf.Redis.MessageRedis.URI) 56 | 57 | ft := messageRedis.NewBloomFilter(config.DatabaseName) 58 | 59 | ope := operations.NewDatabaseOperations(ctx, db, ft, es, ndb) 60 | 61 | tracer := tasktracer.NewTaskTracer(messageRedis, config.TaskTracerChannel) 62 | 63 | resp := responser.NewResponser(conf.Responser.Listen, ope, mq, config.CrawlTaskQueue, conf.General.SpiderRetry, tracer, cacheRedis, conf.Responser.Debug, conf.Responser.Mode, conf.Responser.EnableForceFetch) 64 | 65 | err = resp.Run() 66 | if err != nil { 67 | log.Fatal(err) 68 | } 69 | } 70 | -------------------------------------------------------------------------------- /spider.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 6 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 7 | "github.com/ShugetsuSoft/pixivel-back/modules" 8 | "github.com/ShugetsuSoft/pixivel-back/modules/spider" 9 | "log" 10 | ) 11 | 12 | func main() { 13 | var configPath string 14 | flag.StringVar(&configPath, "config", "config.yaml", "Config File Path") 15 | flag.Parse() 16 | 17 | conf, err := modules.ReadConfig(configPath) 18 | if err != nil { 19 | log.Fatal(err) 20 | } 21 | 22 | telemetry.RegisterSpider() 23 | errchan := telemetry.RunLoki(conf.General.Loki, "spider") 24 | go func() { 25 | for err := range errchan { 26 | log.Println(err) 27 | } 28 | }() 29 | go telemetry.RunPrometheus(conf.General.Prometheus) 30 | 31 | mq, err := modules.NewMQ(conf) 32 | if err != nil { 33 | log.Fatal(err) 34 | } 35 | 36 | spiderI, err := spider.NewSpider(mq, config.CrawlTaskQueue, config.CrawlResponsesExchange, conf.Spider.PixivToken, conf.Spider.CrawlingThreads) 37 | if err != nil { 38 | log.Fatal(err) 39 | } 40 | 41 | spiderI.Crawl() 42 | } 43 | -------------------------------------------------------------------------------- /storer.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 7 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 9 | "github.com/ShugetsuSoft/pixivel-back/common/utils/telemetry" 10 | "github.com/ShugetsuSoft/pixivel-back/modules" 11 | "github.com/ShugetsuSoft/pixivel-back/modules/storer" 12 | "log" 13 | ) 14 | 15 | func main() { 16 | ctx, cancel := context.WithCancel(context.Background()) 17 | defer cancel() 18 | 19 | var configPath string 20 | var runType int 21 | flag.StringVar(&configPath, "config", "config.yaml", "Config File Path") 22 | flag.IntVar(&runType, "run", 0, "Store Type") 23 | flag.Parse() 24 | 25 | conf, err := modules.ReadConfig(configPath) 26 | if err != nil { 27 | log.Fatal(err) 28 | } 29 | 30 | telemetry.RegisterStorer() 31 | errchan := telemetry.RunLoki(conf.General.Loki, "responser") 32 | go func() { 33 | for err := range errchan { 34 | log.Println(err) 35 | } 36 | }() 37 | go telemetry.RunPrometheus(conf.General.Prometheus) 38 | 39 | db, err := modules.NewDB(conf, ctx) 40 | if err != nil { 41 | log.Fatal(err) 42 | } 43 | mq, err := modules.NewMQ(conf) 44 | if err != nil { 45 | log.Fatal(err) 46 | } 47 | ndb, err := modules.NewNB(conf, ctx) 48 | if err != nil { 49 | log.Fatal(err) 50 | } 51 | es, err := modules.NewES(conf, ctx) 52 | if err != nil { 53 | log.Fatal(err) 54 | } 55 | 56 | messageRedis := modules.NewRD(conf.Redis.MessageRedis.URI) 57 | tracer := tasktracer.NewTaskTracer(messageRedis, config.TaskTracerChannel) 58 | ft := messageRedis.NewBloomFilter(config.DatabaseName) 59 | 60 | switch runType { 61 | case 0: 62 | ope := operations.NewDatabaseOperations(ctx, db, ft, es, ndb) 63 | stor := storer.NewStorer(mq, config.CrawlTaskQueue, conf.General.SpiderRetry, ope, tracer) 64 | stor.StoreDB(config.CrawlResponsesMongodb) 65 | case 1: 66 | ope := operations.NewDatabaseOperations(ctx, db, ft, es, ndb) 67 | stor := storer.NewStorer(mq, config.CrawlTaskQueue, conf.General.SpiderRetry, ope, tracer) 68 | stor.StoreES(config.CrawlResponsesElastic) 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /tasks.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 7 | "github.com/ShugetsuSoft/pixivel-back/common/database/tasktracer" 8 | "github.com/ShugetsuSoft/pixivel-back/common/utils/config" 9 | "github.com/ShugetsuSoft/pixivel-back/modules" 10 | "github.com/ShugetsuSoft/pixivel-back/modules/responser/task" 11 | "github.com/ShugetsuSoft/pixivel-back/tools/incrementor" 12 | "log" 13 | "time" 14 | ) 15 | 16 | func main() { 17 | ctx, cancel := context.WithCancel(context.Background()) 18 | defer cancel() 19 | 20 | var configPath string 21 | var date string 22 | flag.StringVar(&configPath, "config", "config.yaml", "Config File Path") 23 | flag.StringVar(&date, "date", "", "Specify rank date") 24 | 25 | flag.Parse() 26 | 27 | conf, err := modules.ReadConfig(configPath) 28 | if err != nil { 29 | log.Fatal(err) 30 | } 31 | 32 | db, err := modules.NewDB(conf, ctx) 33 | if err != nil { 34 | log.Fatal(err) 35 | } 36 | mq, err := modules.NewMQ(conf) 37 | if err != nil { 38 | log.Fatal(err) 39 | } 40 | ndb, err := modules.NewNB(conf, ctx) 41 | if err != nil { 42 | log.Fatal(err) 43 | } 44 | es, err := modules.NewES(conf, ctx) 45 | if err != nil { 46 | log.Fatal(err) 47 | } 48 | 49 | messageRedis := modules.NewRD(conf.Redis.MessageRedis.URI) 50 | tracer := tasktracer.NewTaskTracer(messageRedis, config.TaskTracerChannel) 51 | 52 | ft := messageRedis.NewBloomFilter(config.DatabaseName) 53 | ope := operations.NewDatabaseOperations(ctx, db, ft, es, ndb) 54 | 55 | taskgen := task.NewTaskGenerator(mq, config.CrawlTaskQueue, conf.General.SpiderRetry, tracer) 56 | var dateobj time.Time 57 | if date != "" { 58 | dateobj = time.Now().AddDate(0, 0, -2) 59 | } else { 60 | dateobj, err = time.Parse("20060102", date) 61 | if err != nil { 62 | log.Fatal(err) 63 | } 64 | } 65 | err = incrementor.CrawlRank(taskgen, ope, dateobj) 66 | if err != nil { 67 | log.Fatal(err) 68 | } 69 | } 70 | -------------------------------------------------------------------------------- /tools/incrementor/incrementor.go: -------------------------------------------------------------------------------- 1 | package incrementor 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/ShugetsuSoft/pixivel-back/common/database/operations" 8 | "github.com/ShugetsuSoft/pixivel-back/modules/responser/task" 9 | ) 10 | 11 | func CrawlRank(taskgen *task.TaskGenerator, ope *operations.DatabaseOperations, date time.Time) error { 12 | ctx := context.Background() 13 | 14 | contents := map[string][]string{ 15 | "all": {"daily", "weekly", "monthly", "rookie", "original", "male", "female"}, 16 | "illust": {"daily", "weekly", "monthly", "rookie"}, 17 | "manga": {"daily", "weekly", "monthly", "rookie"}, 18 | "ugoira": {"daily", "weekly"}, 19 | } 20 | today := date.Format("20060102") 21 | for content := range contents { 22 | for index := range contents[content] { 23 | exist, err := ope.IsRankExist(ctx, contents[content][index], today, content) 24 | if err != nil { 25 | return err 26 | } 27 | if exist { 28 | continue 29 | } 30 | err = taskgen.RankInitTask(ctx, contents[content][index], today, content) 31 | if err != nil { 32 | return err 33 | } 34 | time.Sleep(time.Second * 30) 35 | } 36 | } 37 | return nil 38 | } 39 | --------------------------------------------------------------------------------