From ff63084014fe8470e546ecee9fd18f0fb704e705 Mon Sep 17 00:00:00 2001 From: Simon Ding Date: Mon, 19 Aug 2024 17:39:37 +0800 Subject: [PATCH] feat: better jav search support --- README.md | 1 + db/db.go | 1 + ent/media.go | 15 ++++++++- ent/media/media.go | 3 ++ ent/media/where.go | 10 ++++++ ent/media_create.go | 18 ++++++++++ ent/media_update.go | 52 +++++++++++++++++++++++++++++ ent/migrate/schema.go | 1 + ent/mutation.go | 75 +++++++++++++++++++++++++++++++++++++++++- ent/schema/media.go | 7 ++++ pkg/tmdb/tmdb.go | 5 +++ server/core/torrent.go | 17 +++++++--- server/watchlist.go | 38 +++++++++++++++++++-- 13 files changed, 235 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 40d3dec..fc5006e 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ Polaris 是一个电视剧和电影的追踪软件。配置好了之后，当剧 - [ ] qbittorrent客户端支持 - [ ] 更多通知客户端支持 - [ ] 第三方watchlist导入支持 +- [ ] 手机客户端 ## 原理 diff --git a/db/db.go b/db/db.go index 78c9234..99ee487 100644 --- a/db/db.go +++ b/db/db.go @@ -150,6 +150,7 @@ func (c *Client) AddMediaWatchlist(m *ent.Media, episodes []int) (*ent.Media, er SetTargetDir(m.TargetDir). SetDownloadHistoryEpisodes(m.DownloadHistoryEpisodes). SetLimiter(m.Limiter). + SetExtras(m.Extras). AddEpisodeIDs(episodes...). Save(context.TODO()) return r, err diff --git a/ent/media.go b/ent/media.go index 1873d72..e2ecb7a 100644 --- a/ent/media.go +++ b/ent/media.go @@ -47,6 +47,8 @@ type Media struct { DownloadHistoryEpisodes bool `json:"download_history_episodes,omitempty"` // Limiter holds the value of the "limiter" field. Limiter schema.MediaLimiter `json:"limiter,omitempty"` + // Extras holds the value of the "extras" field. + Extras schema.MediaExtras `json:"extras,omitempty"` // Edges holds the relations/edges for other nodes in the graph. // The values are being populated by the MediaQuery when eager-loading is set. Edges MediaEdges `json:"edges"` @@ -76,7 +78,7 @@ func (*Media) scanValues(columns []string) ([]any, error) { values := make([]any, len(columns)) for i := range columns { switch columns[i] { - case media.FieldLimiter: + case media.FieldLimiter, media.FieldExtras: values[i] = new([]byte) case media.FieldDownloadHistoryEpisodes: values[i] = new(sql.NullBool) @@ -193,6 +195,14 @@ func (m *Media) assignValues(columns []string, values []any) error { return fmt.Errorf("unmarshal field limiter: %w", err) } } + case media.FieldExtras: + if value, ok := values[i].(*[]byte); !ok { + return fmt.Errorf("unexpected type %T for field extras", values[i]) + } else if value != nil && len(*value) > 0 { + if err := json.Unmarshal(*value, &m.Extras); err != nil { + return fmt.Errorf("unmarshal field extras: %w", err) + } + } default: m.selectValues.Set(columns[i], values[i]) } @@ -275,6 +285,9 @@ func (m *Media) String() string { builder.WriteString(", ") builder.WriteString("limiter=") builder.WriteString(fmt.Sprintf("%v", m.Limiter)) + builder.WriteString(", ") + builder.WriteString("extras=") + builder.WriteString(fmt.Sprintf("%v", m.Extras)) builder.WriteByte(')') return builder.String() } diff --git a/ent/media/media.go b/ent/media/media.go index 8b0b13d..e092455 100644 --- a/ent/media/media.go +++ b/ent/media/media.go @@ -43,6 +43,8 @@ const ( FieldDownloadHistoryEpisodes = "download_history_episodes" // FieldLimiter holds the string denoting the limiter field in the database. FieldLimiter = "limiter" + // FieldExtras holds the string denoting the extras field in the database. + FieldExtras = "extras" // EdgeEpisodes holds the string denoting the episodes edge name in mutations. EdgeEpisodes = "episodes" // Table holds the table name of the media in the database. @@ -73,6 +75,7 @@ var Columns = []string{ FieldTargetDir, FieldDownloadHistoryEpisodes, FieldLimiter, + FieldExtras, } // ValidColumn reports if the column name is valid (part of the table columns). diff --git a/ent/media/where.go b/ent/media/where.go index d51232c..7652745 100644 --- a/ent/media/where.go +++ b/ent/media/where.go @@ -785,6 +785,16 @@ func LimiterNotNil() predicate.Media { return predicate.Media(sql.FieldNotNull(FieldLimiter)) } +// ExtrasIsNil applies the IsNil predicate on the "extras" field. +func ExtrasIsNil() predicate.Media { + return predicate.Media(sql.FieldIsNull(FieldExtras)) +} + +// ExtrasNotNil applies the NotNil predicate on the "extras" field. +func ExtrasNotNil() predicate.Media { + return predicate.Media(sql.FieldNotNull(FieldExtras)) +} + // HasEpisodes applies the HasEdge predicate on the "episodes" edge. func HasEpisodes() predicate.Media { return predicate.Media(func(s *sql.Selector) { diff --git a/ent/media_create.go b/ent/media_create.go index e6ae0eb..cacf0e8 100644 --- a/ent/media_create.go +++ b/ent/media_create.go @@ -170,6 +170,20 @@ func (mc *MediaCreate) SetNillableLimiter(sl *schema.MediaLimiter) *MediaCreate return mc } +// SetExtras sets the "extras" field. +func (mc *MediaCreate) SetExtras(se schema.MediaExtras) *MediaCreate { + mc.mutation.SetExtras(se) + return mc +} + +// SetNillableExtras sets the "extras" field if the given value is not nil. +func (mc *MediaCreate) SetNillableExtras(se *schema.MediaExtras) *MediaCreate { + if se != nil { + mc.SetExtras(*se) + } + return mc +} + // AddEpisodeIDs adds the "episodes" edge to the Episode entity by IDs. func (mc *MediaCreate) AddEpisodeIDs(ids ...int) *MediaCreate { mc.mutation.AddEpisodeIDs(ids...) @@ -359,6 +373,10 @@ func (mc *MediaCreate) createSpec() (*Media, *sqlgraph.CreateSpec) { _spec.SetField(media.FieldLimiter, field.TypeJSON, value) _node.Limiter = value } + if value, ok := mc.mutation.Extras(); ok { + _spec.SetField(media.FieldExtras, field.TypeJSON, value) + _node.Extras = value + } if nodes := mc.mutation.EpisodesIDs(); len(nodes) > 0 { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.O2M, diff --git a/ent/media_update.go b/ent/media_update.go index 65fcb51..a7c0348 100644 --- a/ent/media_update.go +++ b/ent/media_update.go @@ -270,6 +270,26 @@ func (mu *MediaUpdate) ClearLimiter() *MediaUpdate { return mu } +// SetExtras sets the "extras" field. +func (mu *MediaUpdate) SetExtras(se schema.MediaExtras) *MediaUpdate { + mu.mutation.SetExtras(se) + return mu +} + +// SetNillableExtras sets the "extras" field if the given value is not nil. +func (mu *MediaUpdate) SetNillableExtras(se *schema.MediaExtras) *MediaUpdate { + if se != nil { + mu.SetExtras(*se) + } + return mu +} + +// ClearExtras clears the value of the "extras" field. +func (mu *MediaUpdate) ClearExtras() *MediaUpdate { + mu.mutation.ClearExtras() + return mu +} + // AddEpisodeIDs adds the "episodes" edge to the Episode entity by IDs. func (mu *MediaUpdate) AddEpisodeIDs(ids ...int) *MediaUpdate { mu.mutation.AddEpisodeIDs(ids...) @@ -428,6 +448,12 @@ func (mu *MediaUpdate) sqlSave(ctx context.Context) (n int, err error) { if mu.mutation.LimiterCleared() { _spec.ClearField(media.FieldLimiter, field.TypeJSON) } + if value, ok := mu.mutation.Extras(); ok { + _spec.SetField(media.FieldExtras, field.TypeJSON, value) + } + if mu.mutation.ExtrasCleared() { + _spec.ClearField(media.FieldExtras, field.TypeJSON) + } if mu.mutation.EpisodesCleared() { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.O2M, @@ -733,6 +759,26 @@ func (muo *MediaUpdateOne) ClearLimiter() *MediaUpdateOne { return muo } +// SetExtras sets the "extras" field. +func (muo *MediaUpdateOne) SetExtras(se schema.MediaExtras) *MediaUpdateOne { + muo.mutation.SetExtras(se) + return muo +} + +// SetNillableExtras sets the "extras" field if the given value is not nil. +func (muo *MediaUpdateOne) SetNillableExtras(se *schema.MediaExtras) *MediaUpdateOne { + if se != nil { + muo.SetExtras(*se) + } + return muo +} + +// ClearExtras clears the value of the "extras" field. +func (muo *MediaUpdateOne) ClearExtras() *MediaUpdateOne { + muo.mutation.ClearExtras() + return muo +} + // AddEpisodeIDs adds the "episodes" edge to the Episode entity by IDs. func (muo *MediaUpdateOne) AddEpisodeIDs(ids ...int) *MediaUpdateOne { muo.mutation.AddEpisodeIDs(ids...) @@ -921,6 +967,12 @@ func (muo *MediaUpdateOne) sqlSave(ctx context.Context) (_node *Media, err error if muo.mutation.LimiterCleared() { _spec.ClearField(media.FieldLimiter, field.TypeJSON) } + if value, ok := muo.mutation.Extras(); ok { + _spec.SetField(media.FieldExtras, field.TypeJSON, value) + } + if muo.mutation.ExtrasCleared() { + _spec.ClearField(media.FieldExtras, field.TypeJSON) + } if muo.mutation.EpisodesCleared() { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.O2M, diff --git a/ent/migrate/schema.go b/ent/migrate/schema.go index aa85afd..04be080 100644 --- a/ent/migrate/schema.go +++ b/ent/migrate/schema.go @@ -110,6 +110,7 @@ var ( {Name: "target_dir", Type: field.TypeString, Nullable: true}, {Name: "download_history_episodes", Type: field.TypeBool, Nullable: true, Default: false}, {Name: "limiter", Type: field.TypeJSON, Nullable: true}, + {Name: "extras", Type: field.TypeJSON, Nullable: true}, } // MediaTable holds the schema information for the "media" table. MediaTable = &schema.Table{ diff --git a/ent/mutation.go b/ent/mutation.go index d51c9d5..bcc9fb8 100644 --- a/ent/mutation.go +++ b/ent/mutation.go @@ -3675,6 +3675,7 @@ type MediaMutation struct { target_dir *string download_history_episodes *bool limiter *schema.MediaLimiter + extras *schema.MediaExtras clearedFields map[string]struct{} episodes map[int]struct{} removedepisodes map[int]struct{} @@ -4392,6 +4393,55 @@ func (m *MediaMutation) ResetLimiter() { delete(m.clearedFields, media.FieldLimiter) } +// SetExtras sets the "extras" field. +func (m *MediaMutation) SetExtras(se schema.MediaExtras) { + m.extras = &se +} + +// Extras returns the value of the "extras" field in the mutation. +func (m *MediaMutation) Extras() (r schema.MediaExtras, exists bool) { + v := m.extras + if v == nil { + return + } + return *v, true +} + +// OldExtras returns the old "extras" field's value of the Media entity. +// If the Media object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *MediaMutation) OldExtras(ctx context.Context) (v schema.MediaExtras, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldExtras is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldExtras requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldExtras: %w", err) + } + return oldValue.Extras, nil +} + +// ClearExtras clears the value of the "extras" field. +func (m *MediaMutation) ClearExtras() { + m.extras = nil + m.clearedFields[media.FieldExtras] = struct{}{} +} + +// ExtrasCleared returns if the "extras" field was cleared in this mutation. +func (m *MediaMutation) ExtrasCleared() bool { + _, ok := m.clearedFields[media.FieldExtras] + return ok +} + +// ResetExtras resets all changes to the "extras" field. +func (m *MediaMutation) ResetExtras() { + m.extras = nil + delete(m.clearedFields, media.FieldExtras) +} + // AddEpisodeIDs adds the "episodes" edge to the Episode entity by ids. func (m *MediaMutation) AddEpisodeIDs(ids ...int) { if m.episodes == nil { @@ -4480,7 +4530,7 @@ func (m *MediaMutation) Type() string { // order to get all numeric fields that were incremented/decremented, call // AddedFields(). func (m *MediaMutation) Fields() []string { - fields := make([]string, 0, 14) + fields := make([]string, 0, 15) if m.tmdb_id != nil { fields = append(fields, media.FieldTmdbID) } @@ -4523,6 +4573,9 @@ func (m *MediaMutation) Fields() []string { if m.limiter != nil { fields = append(fields, media.FieldLimiter) } + if m.extras != nil { + fields = append(fields, media.FieldExtras) + } return fields } @@ -4559,6 +4612,8 @@ func (m *MediaMutation) Field(name string) (ent.Value, bool) { return m.DownloadHistoryEpisodes() case media.FieldLimiter: return m.Limiter() + case media.FieldExtras: + return m.Extras() } return nil, false } @@ -4596,6 +4651,8 @@ func (m *MediaMutation) OldField(ctx context.Context, name string) (ent.Value, e return m.OldDownloadHistoryEpisodes(ctx) case media.FieldLimiter: return m.OldLimiter(ctx) + case media.FieldExtras: + return m.OldExtras(ctx) } return nil, fmt.Errorf("unknown Media field %s", name) } @@ -4703,6 +4760,13 @@ func (m *MediaMutation) SetField(name string, value ent.Value) error { } m.SetLimiter(v) return nil + case media.FieldExtras: + v, ok := value.(schema.MediaExtras) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetExtras(v) + return nil } return fmt.Errorf("unknown Media field %s", name) } @@ -4775,6 +4839,9 @@ func (m *MediaMutation) ClearedFields() []string { if m.FieldCleared(media.FieldLimiter) { fields = append(fields, media.FieldLimiter) } + if m.FieldCleared(media.FieldExtras) { + fields = append(fields, media.FieldExtras) + } return fields } @@ -4804,6 +4871,9 @@ func (m *MediaMutation) ClearField(name string) error { case media.FieldLimiter: m.ClearLimiter() return nil + case media.FieldExtras: + m.ClearExtras() + return nil } return fmt.Errorf("unknown Media nullable field %s", name) } @@ -4854,6 +4924,9 @@ func (m *MediaMutation) ResetField(name string) error { case media.FieldLimiter: m.ResetLimiter() return nil + case media.FieldExtras: + m.ResetExtras() + return nil } return fmt.Errorf("unknown Media field %s", name) } diff --git a/ent/schema/media.go b/ent/schema/media.go index 9ab821d..19fe211 100644 --- a/ent/schema/media.go +++ b/ent/schema/media.go @@ -30,6 +30,7 @@ func (Media) Fields() []ent.Field { field.String("target_dir").Optional(), field.Bool("download_history_episodes").Optional().Default(false).Comment("tv series only"), field.JSON("limiter", MediaLimiter{}).Optional(), + field.JSON("extras", MediaExtras{}).Optional(), } } @@ -44,3 +45,9 @@ type MediaLimiter struct { SizeMin int `json:"size_min"` //in B SizeMax int `json:"size_max"` //in B } + +type MediaExtras struct { + IsAdult bool `json:"is_adult"` + IsJav bool `json:"is_jav"` + JavId string `json:"javid"` +} diff --git a/pkg/tmdb/tmdb.go b/pkg/tmdb/tmdb.go index ef78538..0cfd26e 100644 --- a/pkg/tmdb/tmdb.go +++ b/pkg/tmdb/tmdb.go @@ -209,6 +209,11 @@ func (c *Client) GetTVAlternativeTitles(id int, language string) (*tmdb.TVAltern return c.tmdbClient.GetTVAlternativeTitles(id, withLangOption(language)) } +func (c *Client) GetMovieAlternativeTitles(id int, language string) (*tmdb.MovieAlternativeTitles, error) { + return c.tmdbClient.GetMovieAlternativeTitles(id, withLangOption(language)) +} + + func wrapLanguage(lang string) string { if lang == "" { lang = "zh-CN" diff --git a/server/core/torrent.go b/server/core/torrent.go index d0678df..0a165df 100644 --- a/server/core/torrent.go +++ b/server/core/torrent.go @@ -160,6 +160,10 @@ func SearchMovie(db1 *db.Client, param *SearchParam) ([]torznab.Result, error) { } res := searchWithTorznab(db1, movieDetail.NameEn, movieDetail.NameCn, movieDetail.OriginalName) + if movieDetail.Extras.IsJav { + res1 := searchWithTorznab(db1, movieDetail.Extras.JavId) + res = append(res, res1...) + } if len(res) == 0 { return nil, fmt.Errorf("no resource found") @@ -176,10 +180,12 @@ func SearchMovie(db1 *db.Client, param *SearchParam) ([]torznab.Result, error) { if !torrentNameOk(movieDetail, r.Name) { continue } - ss := strings.Split(movieDetail.AirDate, "-")[0] - year, _ := strconv.Atoi(ss) - if meta.Year != year && meta.Year != year-1 && meta.Year != year+1 { //year not match - continue + if !movieDetail.Extras.IsJav { + ss := strings.Split(movieDetail.AirDate, "-")[0] + year, _ := strconv.Atoi(ss) + if meta.Year != year && meta.Year != year-1 && meta.Year != year+1 { //year not match + continue + } } } @@ -295,6 +301,9 @@ func dedup(list []torznab.Result) []torznab.Result { } func torrentNameOk(detail *db.MediaDetails, torrentName string) bool { + if detail.Extras.IsJav && isNameAcceptable(torrentName, detail.Extras.JavId) { + return true + } return isNameAcceptable(torrentName, detail.NameCn) || isNameAcceptable(torrentName, detail.NameEn) || isNameAcceptable(torrentName, detail.OriginalName) } diff --git a/server/watchlist.go b/server/watchlist.go index ff3c647..657094b 100644 --- a/server/watchlist.go +++ b/server/watchlist.go @@ -13,6 +13,7 @@ import ( "polaris/ent/schema" "polaris/log" "strconv" + "strings" "time" tmdb "github.com/cyruzin/golang-tmdb" @@ -174,6 +175,26 @@ func (s *Server) AddTv2Watchlist(c *gin.Context) (interface{}, error) { return nil, nil } +func isJav(detail *tmdb.MovieDetails) bool { + if detail.Adult && len(detail.ProductionCountries)> 0 && strings.ToUpper(detail.ProductionCountries[0].Iso3166_1) == "JP" { + return true + } + return false +} + +func (s *Server) getJavid(id int) string { + alters, err := s.MustTMDB().GetMovieAlternativeTitles(id, s.language) + if err != nil { + return "" + } + for _, t := range alters.Titles { + if t.Iso3166_1 == "JP" && t.Type == "" { + return t.Title + } + } + return "" +} + func (s *Server) AddMovie2Watchlist(c *gin.Context) (interface{}, error) { var in addWatchlistIn if err := c.ShouldBindJSON(&in); err != nil { @@ -196,6 +217,7 @@ func (s *Server) AddMovie2Watchlist(c *gin.Context) (interface{}, error) { } log.Infof("find detail for movie id %d: %v", in.TmdbID, detail) + epid, err := s.db.SaveEposideDetail(&ent.Episode{ SeasonNumber: 1, EpisodeNumber: 1, @@ -209,7 +231,7 @@ func (s *Server) AddMovie2Watchlist(c *gin.Context) (interface{}, error) { } log.Infof("added dummy episode for movie: %v", nameEn) - r, err := s.db.AddMediaWatchlist(&ent.Media{ + movie := ent.Media{ TmdbID: int(detail.ID), ImdbID: detail.IMDbID, MediaType: media.MediaTypeMovie, @@ -222,7 +244,19 @@ func (s *Server) AddMovie2Watchlist(c *gin.Context) (interface{}, error) { StorageID: in.StorageID, TargetDir: in.Folder, Limiter: schema.MediaLimiter{SizeMin: in.SizeMin, SizeMax: in.SizeMax}, - }, []int{epid}) + } + if isJav(detail) { + javid := s.getJavid(in.TmdbID) + movie.Extras = schema.MediaExtras{ + IsJav: true, + JavId: javid, + } + } + if detail.Adult { + movie.Extras.IsAdult = true + } + + r, err := s.db.AddMediaWatchlist(&movie, []int{epid}) if err != nil { return nil, errors.Wrap(err, "add to list") }