flexget.plugins.filter.movie_queue
Covered: 87 lines
Missed: 71 lines
Skipped 51 lines
Percent: 55 %
  1
import logging
  2
from sqlalchemy import Column, Integer, String, ForeignKey, or_, and_
  3
from flexget import schema
  4
from flexget.manager import Session
  5
from flexget.utils import qualities
  6
from flexget.utils.imdb import extract_id
  7
from flexget.utils.database import quality_property, with_session
  8
from flexget.utils.sqlalchemy_utils import table_exists, table_schema
  9
from flexget.plugin import DependencyError, get_plugin_by_name, register_plugin
 10
from flexget.event import event
 12
try:
 13
    from flexget.plugins.filter import queue_base
 14
except ImportError:
 15
    raise DependencyError(issued_by='movie_queue', missing='queue_base',
 16
                             message='movie_queue requires the queue_base plugin')
 18
log = logging.getLogger('movie_queue')
 19
Base = schema.versioned_base('movie_queue', 0)
 22
@event('manager.startup')
 23
def migrate_imdb_queue(manager):
 24
    """If imdb_queue table is found, migrate the data to movie_queue"""
 25
    session = Session()
 26
    try:
 27
        if table_exists('imdb_queue', session):
 28
            log.info('Migrating imdb_queue items to movie_queue')
 29
            old_table = table_schema('imdb_queue', session)
 30
            for row in session.execute(old_table.select()):
 31
                try:
 32
                    queue_add(imdb_id=row['imdb_id'], quality=row['quality'], force=row['immortal'], session=session)
 33
                except QueueError, e:
 34
                    log.error('Unable to migrate %s from imdb_queue to movie_queue' % row['title'])
 35
            old_table.drop()
 36
            session.commit()
 37
    finally:
 38
        session.close()
 41
class QueuedMovie(queue_base.QueuedItem, Base):
 42
    __tablename__ = 'movie_queue'
 43
    __mapper_args__ = {'polymorphic_identity': 'movie'}
 44
    id = Column(Integer, ForeignKey('queue.id'), primary_key=True)
 45
    imdb_id = Column(String)
 46
    tmdb_id = Column(Integer)
 47
    quality = Column('quality', String)
 48
    quality_obj = quality_property('quality')
 51
class FilterMovieQueue(queue_base.FilterQueueBase):
 53
    def matches(self, feed, config, entry):
 55
        try:
 56
            get_plugin_by_name('tmdb_lookup').instance.lookup(entry)
 57
        except DependencyError:
 58
            log.debug('tmdb_lookup is not available, queue will not work if movie ids are not populated')
 60
        conditions = []
 62
        for lazy in [False, True]:
 63
            if entry.get('imdb_id', eval_lazy=lazy):
 64
                conditions.append(QueuedMovie.imdb_id == entry['imdb_id'])
 65
            if entry.get('tmdb_id', eval_lazy=lazy):
 66
                conditions.append(QueuedMovie.tmdb_id == entry['tmdb_id'])
 67
            if conditions:
 68
                break
 69
        if not conditions:
 70
            log.warning('No movie id could be determined for %s' % entry['title'])
 71
            return
 73
        quality = entry.get('quality', qualities.UNKNOWN)
 75
        return feed.session.query(QueuedMovie).filter(QueuedMovie.downloaded == None).\
 76
                                               filter(or_(*conditions)).\
 77
                                               filter(QueuedMovie.quality_obj <= quality).first()
 80
class QueueError(Exception):
 81
    """Exception raised if there is an error with a queue operation"""
 86
    def __init__(self, message, errno=0):
 87
        self.message = message
 88
        self.errno = errno
 91
def validate_quality(quality):
 94
    if isinstance(quality, qualities.Quality):
 95
        if quality.value <= 0:
 96
            return 'ANY'
 97
        return quality.name
 98
    elif quality.upper() == 'ANY':
 99
        return 'ANY'
100
    elif qualities.get(quality, False):
101
        return qualities.common_name(quality)
102
    else:
103
        raise QueueError('ERROR! Unknown quality `%s`' % quality, errno=1)
106
@with_session
107
def parse_what(what, session=None):
108
    """Parses needed movie information for a given search string.
110
    Search string can be one of:
111
        <Movie Title>: Search based on title
112
        imdb_id=<IMDB id>: search based on imdb id
113
        tmdb_id=<TMDB id>: search based on tmdb id"""
115
    tmdb_lookup = get_plugin_by_name('api_tmdb').instance.lookup
117
    imdb_id = extract_id(what)
118
    try:
119
        if imdb_id:
120
            movie = tmdb_lookup(imdb_id=imdb_id, session=session)
121
        elif what.startswith('tmdb_id='):
122
            movie = tmdb_lookup(tmdb_id=what[8:], session=session)
123
        else:
124
            movie = tmdb_lookup(title=what, session=session)
125
    except LookupError, e:
126
        raise QueueError(e.message)
128
    if movie:
129
        return {'title': movie.name, 'imdb_id': movie.imdb_id, 'tmdb_id': movie.id}
130
    else:
131
        raise QueueError('ERROR: Unable to find any such movie from tmdb, use imdb or tmdb id instead.')
135
@with_session
136
def queue_add(title=None, imdb_id=None, tmdb_id=None, quality='ANY', force=True, session=None):
137
    """Add an item to the queue with the specified quality"""
139
    if not title or not (imdb_id or tmdb_id):
141
        result = parse_what(imdb_id or title, session=session)
142
        title = result['title']
143
        imdb_id = result['imdb_id']
144
        tmdb_id = result['tmdb_id']
145
    quality = validate_quality(quality)
148
    item = session.query(QueuedMovie).filter(or_(and_(QueuedMovie.imdb_id != None, QueuedMovie.imdb_id == imdb_id),
149
                                                 and_(QueuedMovie.tmdb_id != None, QueuedMovie.tmdb_id == tmdb_id))).\
150
                                      first()
151
    if not item:
152
        item = QueuedMovie(title=title, imdb_id=imdb_id, tmdb_id=tmdb_id, quality=quality, immortal=force)
153
        session.add(item)
154
        log.info('Adding %s to movie queue with quality=%s and force=%s.' % (title, quality, force))
155
        return {'title': title, 'imdb_id': imdb_id, 'tmdb_id': tmdb_id, 'quality': quality, 'force': force}
156
    else:
157
        if item.downloaded:
158
            raise QueueError('ERROR: %s has already been queued and downloaded' % title)
159
        else:
160
            raise QueueError('ERROR: %s is already in the queue' % title)
163
@with_session
164
def queue_del(imdb_id, session=None):
165
    """Delete the given item from the queue"""
168
    item = session.query(QueuedMovie).filter(QueuedMovie.imdb_id == imdb_id).first()
169
    if item:
170
        title = item.title
171
        session.delete(item)
172
        return title
173
    else:
174
        raise QueueError('%s is not in the queue' % imdb_id)
177
@with_session
178
def queue_edit(imdb_id, quality, session=None):
179
    """Change the required quality for a movie in the queue"""
180
    quality = validate_quality(quality)
182
    item = session.query(QueuedMovie).filter(QueuedMovie.imdb_id == imdb_id).first()
183
    if item:
184
        item.quality = quality
185
        session.commit()
186
        return item.title
187
    else:
188
        raise QueueError('%s is not in the queue' % imdb_id)
191
@with_session
192
def queue_get(session=None, downloaded=False):
193
    """Get the current IMDb queue.
195
    KWArgs:
196
        session: new session is used it not given
197
        downloaded: boolean whether or not to return only downloaded
199
    Returns:
200
        List of QueuedMovie objects (detached from session)
201
    """
202
    if not downloaded:
203
        return session.query(QueuedMovie).filter(QueuedMovie.downloaded == None).all()
204
    else:
205
        return session.query(QueuedMovie).filter(QueuedMovie.downloaded != None).all()
208
register_plugin(FilterMovieQueue, 'movie_queue', api_ver=2)