From b2a9f1f83b4f67f32beaad4da3c0703aaaa37bd9 Mon Sep 17 00:00:00 2001 From: Simon Ding Date: Tue, 19 Nov 2024 19:24:43 +0800 Subject: [PATCH] refactor: size limiter --- db/const.go | 25 ++++++++++++++++++++++--- db/db.go | 33 +++++++++++++++------------------ ent/schema/media.go | 7 ++++--- pkg/torznab/torznab.go | 12 ++++++------ server/core/importlist.go | 11 ++++++----- server/core/resources.go | 6 +++--- server/core/scheduler.go | 2 +- server/core/torrent.go | 10 +++++----- server/setting.go | 30 +++++++++--------------------- 9 files changed, 71 insertions(+), 65 deletions(-) diff --git a/db/const.go b/db/const.go index ae3a79c..a22ac5f 100644 --- a/db/const.go +++ b/db/const.go @@ -1,5 +1,7 @@ package db +import "polaris/ent/media" + var Version = "undefined" const ( @@ -17,9 +19,9 @@ const ( SettingTvNamingFormat = "tv_naming_format" SettingMovieNamingFormat = "movie_naming_format" SettingProwlarrInfo = "prowlarr_info" - Setting720pSizeLimiter = "720p_size_limiter" - Setting1080ppSizeLimiter = "1080p_size_limiter" - Setting2160ppSizeLimiter = "2160p_size_limiter" + + SettingTvSizeLimiter = "tv_size_limiter" + SettingMovieSizeLimiter = "movie_size_limiter" ) const ( @@ -53,6 +55,23 @@ type ResolutionType string const JwtSerectKey = "jwt_secrect_key" +type MediaSizeLimiter struct { + P720p SizeLimiter `json:"720p"` + P1080 SizeLimiter `json:"1080p"` + P2160 SizeLimiter `json:"2160p"` +} + +func (m *MediaSizeLimiter) GetLimiter(r media.Resolution) SizeLimiter { + if r == media.Resolution1080p { + return m.P1080 + } else if r == media.Resolution720p { + return m.P720p + } else if r == media.Resolution2160p { + return m.P2160 + } + return SizeLimiter{} +} + type SizeLimiter struct { MaxSIze int64 `json:"max_size"` MinSize int64 `json:"min_size"` diff --git a/db/db.go b/db/db.go index 6c7c450..8a94684 100644 --- a/db/db.go +++ b/db/db.go @@ -441,7 +441,7 @@ type Storage struct { } func (s *Storage) ToWebDavSetting() WebdavSetting { - if s.Implementation != storage.ImplementationWebdav && s.Implementation != storage.ImplementationAlist{ + if s.Implementation != storage.ImplementationWebdav && s.Implementation != storage.ImplementationAlist { panic("not webdav storage") } var webdavSetting WebdavSetting @@ -631,19 +631,17 @@ func (c *Client) DeleteImportlist(id int) error { return c.ent.ImportList.DeleteOneID(id).Exec(context.TODO()) } -func (c *Client) GetSizeLimiter(r media.Resolution) (*SizeLimiter, error) { +func (c *Client) GetSizeLimiter(mediaType string) (*MediaSizeLimiter, error) { var v string - if r == media.Resolution720p { - v = c.GetSetting(Setting720pSizeLimiter) - } else if r == media.Resolution1080p { - v = c.GetSetting(Setting1080ppSizeLimiter) - } else if r == media.Resolution2160p { - v = c.GetSetting(Setting2160ppSizeLimiter) + if mediaType == "tv" { + v = c.GetSetting(SettingTvSizeLimiter) + } else if mediaType == "movie" { + v = c.GetSetting(SettingMovieSizeLimiter) } else { - return nil, errors.Errorf("resolution not supported: %v", r) + return nil, errors.Errorf("media type not supported: %v", mediaType) } - var limiter SizeLimiter + var limiter MediaSizeLimiter if v == "" { return &limiter, nil } @@ -652,20 +650,19 @@ func (c *Client) GetSizeLimiter(r media.Resolution) (*SizeLimiter, error) { return &limiter, err } -func (c *Client) SetSizeLimiter(r media.Resolution, limiter *SizeLimiter) error { +func (c *Client) SetSizeLimiter(mediaType string, limiter *MediaSizeLimiter) error { data, err := json.Marshal(limiter) if err != nil { return err } - if r == media.Resolution720p { - return c.SetSetting(Setting720pSizeLimiter, string(data)) - } else if r == media.Resolution1080p { - return c.SetSetting(Setting1080ppSizeLimiter, string(data)) - } else if r == media.Resolution2160p { - return c.SetSetting(Setting2160ppSizeLimiter, string(data)) + if mediaType == "tv" { + return c.SetSetting(SettingTvSizeLimiter, string(data)) + } else if mediaType == "movie" { + return c.SetSetting(SettingMovieSizeLimiter, string(data)) } else { - return errors.Errorf("resolution not supported: %v", r) + return errors.Errorf("media type not supported: %v", mediaType) } + } func (c *Client) GetTvNamingFormat() string { diff --git a/ent/schema/media.go b/ent/schema/media.go index bc18b1c..1119c15 100644 --- a/ent/schema/media.go +++ b/ent/schema/media.go @@ -42,8 +42,9 @@ func (Media) Edges() []ent.Edge { } type MediaLimiter struct { - SizeMin int `json:"size_min"` //in B - SizeMax int `json:"size_max"` //in B + SizeMin int64 `json:"size_min"` //in B + SizeMax int64 `json:"size_max"` //in B + PreferSize int64 `json:"prefer_max"` } type MediaExtras struct { @@ -51,7 +52,7 @@ type MediaExtras struct { JavId string `json:"javid"` //OriginCountry []string `json:"origin_country"` OriginalLanguage string `json:"original_language"` - Genres []struct { + Genres []struct { ID int64 `json:"id"` Name string `json:"name"` } `json:"genres"` diff --git a/pkg/torznab/torznab.go b/pkg/torznab/torznab.go index 47bdef0..b0d45a2 100644 --- a/pkg/torznab/torznab.go +++ b/pkg/torznab/torznab.go @@ -96,9 +96,9 @@ func (r *Response) ToResults(indexer *db.TorznabInfo) []Result { Description: item.Description, Link: item.Link, Size: mustAtoI(item.Size), - Seeders: mustAtoI(item.GetAttr("seeders")), - Peers: mustAtoI(item.GetAttr("peers")), - Category: mustAtoI(item.GetAttr("category")), + Seeders: int(mustAtoI(item.GetAttr("seeders"))), + Peers: int(mustAtoI(item.GetAttr("peers"))), + Category: int(mustAtoI(item.GetAttr("category"))), ImdbId: imdb, DownloadVolumeFactor: tryParseFloat(item.GetAttr("downloadvolumefactor")), UploadVolumeFactor: tryParseFloat(item.GetAttr("uploadvolumefactor")), @@ -112,8 +112,8 @@ func (r *Response) ToResults(indexer *db.TorznabInfo) []Result { return res } -func mustAtoI(key string) int { - i, err := strconv.Atoi(key) +func mustAtoI(key string) int64 { + i, err := strconv.ParseInt(key, 10, 64) if err != nil { log.Errorf("must atoi error: %v", err) panic(err) @@ -183,7 +183,7 @@ type Result struct { Name string `json:"name"` Description string `json:"description"` Link string `json:"link"` - Size int `json:"size"` + Size int64 `json:"size"` Seeders int `json:"seeders"` Peers int `json:"peers"` Category int `json:"category"` diff --git a/server/core/importlist.go b/server/core/importlist.go index fbaaf05..1adcad2 100644 --- a/server/core/importlist.go +++ b/server/core/importlist.go @@ -114,8 +114,9 @@ type AddWatchlistIn struct { Resolution string `json:"resolution" binding:"required"` Folder string `json:"folder" binding:"required"` DownloadHistoryEpisodes bool `json:"download_history_episodes"` //for tv - SizeMin int `json:"size_min"` - SizeMax int `json:"size_max"` + SizeMin int64 `json:"size_min"` + SizeMax int64 `json:"size_max"` + PreferSize int64 `json:"prefer_size"` } func (c *Client) AddTv2Watchlist(in AddWatchlistIn) (interface{}, error) { @@ -139,7 +140,7 @@ func (c *Client) AddTv2Watchlist(in AddWatchlistIn) (interface{}, error) { } log.Infof("find detail for tv id %d: %+v", in.TmdbID, detail) - lastSeason := 0 + lastSeason := 0 for _, season := range detail.Seasons { if season.SeasonNumber > lastSeason && season.EpisodeCount > 0 { //如果最新一季已经有剧集信息，则以最新一季为准 lastSeason = season.SeasonNumber @@ -327,8 +328,8 @@ func (c *Client) checkMovieFolder(m *ent.Media) error { return err } - for _,f := range files { - if f.IsDir() || f.Size() < 100 * 1000 * 1000 /* 100M */{ //忽略路径和小于100M的文件 + for _, f := range files { + if f.IsDir() || f.Size() < 100*1000*1000 /* 100M */ { //忽略路径和小于100M的文件 continue } meta := metadata.ParseMovie(f.Name()) diff --git a/server/core/resources.go b/server/core/resources.go index a6cb3f5..9535cd9 100644 --- a/server/core/resources.go +++ b/server/core/resources.go @@ -68,7 +68,7 @@ func (c *Client) DownloadEpisodeTorrent(r1 torznab.Result, seriesId, seasonNum i SourceTitle: r1.Name, TargetDir: dir, Status: history.StatusRunning, - Size: r1.Size, + Size: int(r1.Size), //Saved: torrent.Save(), Link: magnet, DownloadClientID: dlc.ID, @@ -192,7 +192,7 @@ lo: return torrentNames, nil } -func (c *Client) DownloadMovie(m *ent.Media, link, name string, size int, indexerID int) (*string, error) { +func (c *Client) DownloadMovie(m *ent.Media, link, name string, size int64, indexerID int) (*string, error) { trc, dlc, err := c.GetDownloadClient() if err != nil { return nil, errors.Wrap(err, "connect transmission") @@ -219,7 +219,7 @@ func (c *Client) DownloadMovie(m *ent.Media, link, name string, size int, indexe SourceTitle: name, TargetDir: m.TargetDir, Status: history.StatusRunning, - Size: size, + Size: int(size), //Saved: torrent.Save(), Link: magnet, DownloadClientID: dlc.ID, diff --git a/server/core/scheduler.go b/server/core/scheduler.go index ea992ed..35a5025 100644 --- a/server/core/scheduler.go +++ b/server/core/scheduler.go @@ -434,7 +434,7 @@ func (c *Client) downloadMovieSingleEpisode(ep *ent.Episode, targetDir string) ( SourceTitle: r1.Name, TargetDir: targetDir, Status: history.StatusRunning, - Size: r1.Size, + Size: int(r1.Size), //Saved: torrent.Save(), Link: magnet, DownloadClientID: dlc.ID, diff --git a/server/core/torrent.go b/server/core/torrent.go index 2e9e468..e2c6c12 100644 --- a/server/core/torrent.go +++ b/server/core/torrent.go @@ -114,8 +114,8 @@ func imdbIDMatchExact(id1, id2 string) bool { return id1 == id2 } -func torrentSizeOk(detail *db.MediaDetails, torrentSize int, torrentEpisodeNum int, param *SearchParam) bool { - defaultMinSize := 80 * 1000 * 1000 //tv, 80M min +func torrentSizeOk(detail *db.MediaDetails, torrentSize int64, torrentEpisodeNum int, param *SearchParam) bool { + defaultMinSize := int64(80 * 1000 * 1000) //tv, 80M min if detail.MediaType == media.MediaTypeMovie { defaultMinSize = 200 * 1000 * 1000 // movie, 200M min } @@ -134,20 +134,20 @@ func torrentSizeOk(detail *db.MediaDetails, torrentSize int, torrentEpisodeNum i if param.CheckFileSize { //check file size when trigger automatic download if detail.Limiter.SizeMin > 0 { //min size - sizeMin := detail.Limiter.SizeMin * multiplier + sizeMin := detail.Limiter.SizeMin * int64(multiplier) if torrentSize < sizeMin { //比最小要求的大小还要小, min size not qualify return false } } if detail.Limiter.SizeMax > 0 { //max size - sizeMax := detail.Limiter.SizeMax * multiplier + sizeMax := detail.Limiter.SizeMax * int64(multiplier) if torrentSize > sizeMax { //larger than max size wanted, max size not qualify return false } } } - return torrentSize > defaultMinSize*multiplier + return torrentSize > defaultMinSize*int64(multiplier) } func seasonEpisodeCount(detail *db.MediaDetails, seasonNum int) int { diff --git a/server/setting.go b/server/setting.go index 33af7af..a192104 100644 --- a/server/setting.go +++ b/server/setting.go @@ -7,7 +7,6 @@ import ( "polaris/db" "polaris/ent" "polaris/ent/downloadclients" - "polaris/ent/media" "polaris/log" "polaris/pkg/prowlarr" "polaris/pkg/qbittorrent" @@ -331,29 +330,22 @@ func (s *Server) SaveProwlarrSetting(c *gin.Context) (interface{}, error) { } type ResolutionSizeLimiter struct { - P720 db.SizeLimiter `json:"720p"` - P1080 db.SizeLimiter `json:"1080p"` - P2160 db.SizeLimiter `json:"2160p"` + TvLimiter *db.MediaSizeLimiter `json:"tv_limiter"` + MovieLimiter *db.MediaSizeLimiter `json:"movie_limiter"` } func (s *Server) GetSizeLimiter(c *gin.Context) (interface{}, error) { - p720, err := s.db.GetSizeLimiter(media.Resolution720p) + tv, err := s.db.GetSizeLimiter("tv") if err != nil { return nil, errors.Wrap(err, "db") } - p1080, err := s.db.GetSizeLimiter(media.Resolution1080p) - if err != nil { - return nil, errors.Wrap(err, "db") - } - - p2160, err := s.db.GetSizeLimiter(media.Resolution2160p) + movie, err := s.db.GetSizeLimiter("movie") if err != nil { return nil, errors.Wrap(err, "db") } r := ResolutionSizeLimiter{ - P720: *p720, - P1080: *p1080, - P2160: *p2160, + TvLimiter: tv, + MovieLimiter: movie, } return r, nil } @@ -363,15 +355,11 @@ func (s *Server) SetSizeLimiter(c *gin.Context) (interface{}, error) { if err := c.ShouldBindJSON(&in); err != nil { return nil, err } - if err := s.db.SetSizeLimiter(media.Resolution720p, &in.P720); err != nil { + if err := s.db.SetSizeLimiter("tv", in.TvLimiter); err != nil { return nil, errors.Wrap(err, "db") } - if err := s.db.SetSizeLimiter(media.Resolution1080p, &in.P1080); err != nil { - return nil, errors.Wrap(err, "db") - } - - if err := s.db.SetSizeLimiter(media.Resolution2160p, &in.P2160); err != nil { + if err := s.db.SetSizeLimiter("movie", in.MovieLimiter); err != nil { return nil, errors.Wrap(err, "db") } return "success", nil -} \ No newline at end of file +}