refactor: size limiter

This commit is contained in:
Simon Ding
2024-11-19 19:24:43 +08:00
parent b69881d26b
commit b2a9f1f83b
9 changed files with 71 additions and 65 deletions

View File

@@ -1,5 +1,7 @@
package db package db
import "polaris/ent/media"
var Version = "undefined" var Version = "undefined"
const ( const (
@@ -17,9 +19,9 @@ const (
SettingTvNamingFormat = "tv_naming_format" SettingTvNamingFormat = "tv_naming_format"
SettingMovieNamingFormat = "movie_naming_format" SettingMovieNamingFormat = "movie_naming_format"
SettingProwlarrInfo = "prowlarr_info" SettingProwlarrInfo = "prowlarr_info"
Setting720pSizeLimiter = "720p_size_limiter"
Setting1080ppSizeLimiter = "1080p_size_limiter" SettingTvSizeLimiter = "tv_size_limiter"
Setting2160ppSizeLimiter = "2160p_size_limiter" SettingMovieSizeLimiter = "movie_size_limiter"
) )
const ( const (
@@ -53,6 +55,23 @@ type ResolutionType string
const JwtSerectKey = "jwt_secrect_key" 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 { type SizeLimiter struct {
MaxSIze int64 `json:"max_size"` MaxSIze int64 `json:"max_size"`
MinSize int64 `json:"min_size"` MinSize int64 `json:"min_size"`

View File

@@ -441,7 +441,7 @@ type Storage struct {
} }
func (s *Storage) ToWebDavSetting() WebdavSetting { 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") panic("not webdav storage")
} }
var webdavSetting WebdavSetting var webdavSetting WebdavSetting
@@ -631,19 +631,17 @@ func (c *Client) DeleteImportlist(id int) error {
return c.ent.ImportList.DeleteOneID(id).Exec(context.TODO()) 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 var v string
if r == media.Resolution720p { if mediaType == "tv" {
v = c.GetSetting(Setting720pSizeLimiter) v = c.GetSetting(SettingTvSizeLimiter)
} else if r == media.Resolution1080p { } else if mediaType == "movie" {
v = c.GetSetting(Setting1080ppSizeLimiter) v = c.GetSetting(SettingMovieSizeLimiter)
} else if r == media.Resolution2160p {
v = c.GetSetting(Setting2160ppSizeLimiter)
} else { } 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 == "" { if v == "" {
return &limiter, nil return &limiter, nil
} }
@@ -652,20 +650,19 @@ func (c *Client) GetSizeLimiter(r media.Resolution) (*SizeLimiter, error) {
return &limiter, err 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) data, err := json.Marshal(limiter)
if err != nil { if err != nil {
return err return err
} }
if r == media.Resolution720p { if mediaType == "tv" {
return c.SetSetting(Setting720pSizeLimiter, string(data)) return c.SetSetting(SettingTvSizeLimiter, string(data))
} else if r == media.Resolution1080p { } else if mediaType == "movie" {
return c.SetSetting(Setting1080ppSizeLimiter, string(data)) return c.SetSetting(SettingMovieSizeLimiter, string(data))
} else if r == media.Resolution2160p {
return c.SetSetting(Setting2160ppSizeLimiter, string(data))
} else { } else {
return errors.Errorf("resolution not supported: %v", r) return errors.Errorf("media type not supported: %v", mediaType)
} }
} }
func (c *Client) GetTvNamingFormat() string { func (c *Client) GetTvNamingFormat() string {

View File

@@ -42,8 +42,9 @@ func (Media) Edges() []ent.Edge {
} }
type MediaLimiter struct { type MediaLimiter struct {
SizeMin int `json:"size_min"` //in B SizeMin int64 `json:"size_min"` //in B
SizeMax int `json:"size_max"` //in B SizeMax int64 `json:"size_max"` //in B
PreferSize int64 `json:"prefer_max"`
} }
type MediaExtras struct { type MediaExtras struct {

View File

@@ -96,9 +96,9 @@ func (r *Response) ToResults(indexer *db.TorznabInfo) []Result {
Description: item.Description, Description: item.Description,
Link: item.Link, Link: item.Link,
Size: mustAtoI(item.Size), Size: mustAtoI(item.Size),
Seeders: mustAtoI(item.GetAttr("seeders")), Seeders: int(mustAtoI(item.GetAttr("seeders"))),
Peers: mustAtoI(item.GetAttr("peers")), Peers: int(mustAtoI(item.GetAttr("peers"))),
Category: mustAtoI(item.GetAttr("category")), Category: int(mustAtoI(item.GetAttr("category"))),
ImdbId: imdb, ImdbId: imdb,
DownloadVolumeFactor: tryParseFloat(item.GetAttr("downloadvolumefactor")), DownloadVolumeFactor: tryParseFloat(item.GetAttr("downloadvolumefactor")),
UploadVolumeFactor: tryParseFloat(item.GetAttr("uploadvolumefactor")), UploadVolumeFactor: tryParseFloat(item.GetAttr("uploadvolumefactor")),
@@ -112,8 +112,8 @@ func (r *Response) ToResults(indexer *db.TorznabInfo) []Result {
return res return res
} }
func mustAtoI(key string) int { func mustAtoI(key string) int64 {
i, err := strconv.Atoi(key) i, err := strconv.ParseInt(key, 10, 64)
if err != nil { if err != nil {
log.Errorf("must atoi error: %v", err) log.Errorf("must atoi error: %v", err)
panic(err) panic(err)
@@ -183,7 +183,7 @@ type Result struct {
Name string `json:"name"` Name string `json:"name"`
Description string `json:"description"` Description string `json:"description"`
Link string `json:"link"` Link string `json:"link"`
Size int `json:"size"` Size int64 `json:"size"`
Seeders int `json:"seeders"` Seeders int `json:"seeders"`
Peers int `json:"peers"` Peers int `json:"peers"`
Category int `json:"category"` Category int `json:"category"`

View File

@@ -114,8 +114,9 @@ type AddWatchlistIn struct {
Resolution string `json:"resolution" binding:"required"` Resolution string `json:"resolution" binding:"required"`
Folder string `json:"folder" binding:"required"` Folder string `json:"folder" binding:"required"`
DownloadHistoryEpisodes bool `json:"download_history_episodes"` //for tv DownloadHistoryEpisodes bool `json:"download_history_episodes"` //for tv
SizeMin int `json:"size_min"` SizeMin int64 `json:"size_min"`
SizeMax int `json:"size_max"` SizeMax int64 `json:"size_max"`
PreferSize int64 `json:"prefer_size"`
} }
func (c *Client) AddTv2Watchlist(in AddWatchlistIn) (interface{}, error) { func (c *Client) AddTv2Watchlist(in AddWatchlistIn) (interface{}, error) {
@@ -327,8 +328,8 @@ func (c *Client) checkMovieFolder(m *ent.Media) error {
return err return err
} }
for _,f := range files { for _, f := range files {
if f.IsDir() || f.Size() < 100 * 1000 * 1000 /* 100M */{ //忽略路径和小于100M的文件 if f.IsDir() || f.Size() < 100*1000*1000 /* 100M */ { //忽略路径和小于100M的文件
continue continue
} }
meta := metadata.ParseMovie(f.Name()) meta := metadata.ParseMovie(f.Name())

View File

@@ -68,7 +68,7 @@ func (c *Client) DownloadEpisodeTorrent(r1 torznab.Result, seriesId, seasonNum i
SourceTitle: r1.Name, SourceTitle: r1.Name,
TargetDir: dir, TargetDir: dir,
Status: history.StatusRunning, Status: history.StatusRunning,
Size: r1.Size, Size: int(r1.Size),
//Saved: torrent.Save(), //Saved: torrent.Save(),
Link: magnet, Link: magnet,
DownloadClientID: dlc.ID, DownloadClientID: dlc.ID,
@@ -192,7 +192,7 @@ lo:
return torrentNames, nil 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() trc, dlc, err := c.GetDownloadClient()
if err != nil { if err != nil {
return nil, errors.Wrap(err, "connect transmission") 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, SourceTitle: name,
TargetDir: m.TargetDir, TargetDir: m.TargetDir,
Status: history.StatusRunning, Status: history.StatusRunning,
Size: size, Size: int(size),
//Saved: torrent.Save(), //Saved: torrent.Save(),
Link: magnet, Link: magnet,
DownloadClientID: dlc.ID, DownloadClientID: dlc.ID,

View File

@@ -434,7 +434,7 @@ func (c *Client) downloadMovieSingleEpisode(ep *ent.Episode, targetDir string) (
SourceTitle: r1.Name, SourceTitle: r1.Name,
TargetDir: targetDir, TargetDir: targetDir,
Status: history.StatusRunning, Status: history.StatusRunning,
Size: r1.Size, Size: int(r1.Size),
//Saved: torrent.Save(), //Saved: torrent.Save(),
Link: magnet, Link: magnet,
DownloadClientID: dlc.ID, DownloadClientID: dlc.ID,

View File

@@ -114,8 +114,8 @@ func imdbIDMatchExact(id1, id2 string) bool {
return id1 == id2 return id1 == id2
} }
func torrentSizeOk(detail *db.MediaDetails, torrentSize int, torrentEpisodeNum int, param *SearchParam) bool { func torrentSizeOk(detail *db.MediaDetails, torrentSize int64, torrentEpisodeNum int, param *SearchParam) bool {
defaultMinSize := 80 * 1000 * 1000 //tv, 80M min defaultMinSize := int64(80 * 1000 * 1000) //tv, 80M min
if detail.MediaType == media.MediaTypeMovie { if detail.MediaType == media.MediaTypeMovie {
defaultMinSize = 200 * 1000 * 1000 // movie, 200M min 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 param.CheckFileSize { //check file size when trigger automatic download
if detail.Limiter.SizeMin > 0 { //min size 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 if torrentSize < sizeMin { //比最小要求的大小还要小, min size not qualify
return false return false
} }
} }
if detail.Limiter.SizeMax > 0 { //max size 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 if torrentSize > sizeMax { //larger than max size wanted, max size not qualify
return false return false
} }
} }
} }
return torrentSize > defaultMinSize*multiplier return torrentSize > defaultMinSize*int64(multiplier)
} }
func seasonEpisodeCount(detail *db.MediaDetails, seasonNum int) int { func seasonEpisodeCount(detail *db.MediaDetails, seasonNum int) int {

View File

@@ -7,7 +7,6 @@ import (
"polaris/db" "polaris/db"
"polaris/ent" "polaris/ent"
"polaris/ent/downloadclients" "polaris/ent/downloadclients"
"polaris/ent/media"
"polaris/log" "polaris/log"
"polaris/pkg/prowlarr" "polaris/pkg/prowlarr"
"polaris/pkg/qbittorrent" "polaris/pkg/qbittorrent"
@@ -331,29 +330,22 @@ func (s *Server) SaveProwlarrSetting(c *gin.Context) (interface{}, error) {
} }
type ResolutionSizeLimiter struct { type ResolutionSizeLimiter struct {
P720 db.SizeLimiter `json:"720p"` TvLimiter *db.MediaSizeLimiter `json:"tv_limiter"`
P1080 db.SizeLimiter `json:"1080p"` MovieLimiter *db.MediaSizeLimiter `json:"movie_limiter"`
P2160 db.SizeLimiter `json:"2160p"`
} }
func (s *Server) GetSizeLimiter(c *gin.Context) (interface{}, error) { 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 { if err != nil {
return nil, errors.Wrap(err, "db") return nil, errors.Wrap(err, "db")
} }
p1080, err := s.db.GetSizeLimiter(media.Resolution1080p) movie, err := s.db.GetSizeLimiter("movie")
if err != nil {
return nil, errors.Wrap(err, "db")
}
p2160, err := s.db.GetSizeLimiter(media.Resolution2160p)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "db") return nil, errors.Wrap(err, "db")
} }
r := ResolutionSizeLimiter{ r := ResolutionSizeLimiter{
P720: *p720, TvLimiter: tv,
P1080: *p1080, MovieLimiter: movie,
P2160: *p2160,
} }
return r, nil return r, nil
} }
@@ -363,14 +355,10 @@ func (s *Server) SetSizeLimiter(c *gin.Context) (interface{}, error) {
if err := c.ShouldBindJSON(&in); err != nil { if err := c.ShouldBindJSON(&in); err != nil {
return nil, err 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") return nil, errors.Wrap(err, "db")
} }
if err := s.db.SetSizeLimiter(media.Resolution1080p, &in.P1080); err != nil { if err := s.db.SetSizeLimiter("movie", in.MovieLimiter); err != nil {
return nil, errors.Wrap(err, "db")
}
if err := s.db.SetSizeLimiter(media.Resolution2160p, &in.P2160); err != nil {
return nil, errors.Wrap(err, "db") return nil, errors.Wrap(err, "db")
} }
return "success", nil return "success", nil