media/cache.py

   1 import logging
   2 import os
   3
   4 import json
   5 import threading
   6
   7 from media.item import item_builders, item_id_generators, dict_to_item
   8 import media.file
   9 import media.url
  10 import media.url_from_playlist
  11 import media.radio
  12 from database import MusicDatabase
  13 import variables as var
  14 import util
  15
  16
  17 class MusicCache(dict):
  18     def __init__(self, db: MusicDatabase):
  19         super().__init__()
  20         self.db = db
  21         self.log = logging.getLogger("bot")
  22         self.dir = None
  23         self.files = []
  24         self.file_id_lookup = {}
  25         self.dir_lock = threading.Lock()
  26
  27     def get_item_by_id(self, bot, id):  # Why all these functions need a bot? Because it need the bot to send message!
  28         if id in self:
  29             return self[id]
  30
  31         # if not cached, query the database
  32         item = self.fetch(bot, id)
  33         if item is not None:
  34             self[id] = item
  35             self.log.debug("library: music found in database: %s" % item.format_debug_string())
  36             return item
  37         else:
  38             return None
  39             # print(id)
  40             # raise KeyError("Unable to fetch item from the database! Please try to refresh the cache by !recache.")
  41
  42     def get_item(self, bot, **kwargs):
  43         # kwargs should provide type and id, and parameters to build the item if not existed in the library.
  44         # if cached
  45         if 'id' in kwargs:
  46             id = kwargs['id']
  47         else:
  48             id = item_id_generators[kwargs['type']](**kwargs)
  49
  50         if id in self:
  51             return self[id]
  52
  53         # if not cached, query the database
  54         item = self.fetch(bot, id)
  55         if item is not None:
  56             self[id] = item
  57             self.log.debug("library: music found in database: %s" % item.format_debug_string())
  58             return item
  59
  60         # if not in the database, build one
  61         self[id] = item_builders[kwargs['type']](bot, **kwargs)  # newly built item will not be saved immediately
  62         return self[id]
  63
  64     def get_items_by_tags(self, bot, tags):
  65         music_dicts = self.db.query_music_by_tags(tags)
  66         items = []
  67         if music_dicts:
  68             for music_dict in music_dicts:
  69                 id = music_dict['id']
  70                 self[id] = dict_to_item(bot, music_dict)
  71                 items.append(self[id])
  72
  73         return items
  74
  75     def fetch(self, bot, id):
  76         music_dicts = self.db.query_music(id=id)
  77         if music_dicts:
  78             music_dict = music_dicts[0]
  79             self[id] = dict_to_item(bot, music_dict)
  80             return self[id]
  81         else:
  82             return None
  83
  84     def save(self, id):
  85         self.log.debug("library: music save into database: %s" % self[id].format_debug_string())
  86         self.db.insert_music(self[id].to_dict())
  87
  88     def free_and_delete(self, id):
  89         item = self.get_item_by_id(None, id)
  90         if item:
  91             self.log.debug("library: DELETE item from the database: %s" % item.format_debug_string())
  92
  93             if item.type == 'file' and item.path in self.file_id_lookup:
  94                 if item.path in self.file_id_lookup:
  95                     del self.file_id_lookup[item.path]
  96                 self.files.remove(item.path)
  97                 self.save_dir_cache()
  98             elif item.type == 'url':
  99                 os.remove(item.path)
 100
 101             if item.id in self:
 102                 del self[item.id]
 103             self.db.delete_music(id=item.id)
 104
 105     def free(self, id):
 106         if id in self:
 107             self.log.debug("library: cache freed for item: %s" % self[id].format_debug_string())
 108             del self[id]
 109
 110     def free_all(self):
 111         self.log.debug("library: all cache freed")
 112         self.clear()
 113
 114     def build_dir_cache(self, bot):
 115         self.dir_lock.acquire()
 116         self.log.info("library: rebuild directory cache")
 117         self.files = []
 118         files = util.get_recursive_file_list_sorted(var.music_folder)
 119         self.dir = util.Dir(var.music_folder)
 120         for file in files:
 121             item = self.fetch(bot, item_id_generators['file'](path=file))
 122             if not item:
 123                 item = item_builders['file'](bot, path=file)
 124                 self.log.debug("library: music save into database: %s" % item.format_debug_string())
 125                 self.db.insert_music(item.to_dict())
 126
 127             self.dir.add_file(file)
 128             self.files.append(file)
 129             self.file_id_lookup[file] = item.id
 130
 131         self.save_dir_cache()
 132         self.dir_lock.release()
 133
 134     def save_dir_cache(self):
 135         var.db.set("dir_cache", "files", json.dumps(self.file_id_lookup))
 136
 137     def load_dir_cache(self, bot):
 138         self.dir_lock.acquire()
 139         self.log.info("library: load directory cache from database")
 140         loaded = json.loads(var.db.get("dir_cache", "files"))
 141         self.files = loaded.keys()
 142         self.file_id_lookup = loaded
 143         self.dir = util.Dir(var.music_folder)
 144         for file, id in loaded.items():
 145             self.dir.add_file(file)
 146         self.dir_lock.release()
 147
 148
 149 class CachedItemWrapper:
 150     def __init__(self, lib, id, type, user):
 151         self.lib = lib
 152         self.id = id
 153         self.user = user
 154         self.type = type
 155         self.log = logging.getLogger("bot")
 156         self.version = 0
 157
 158     def item(self):
 159         return self.lib[self.id]
 160
 161     def to_dict(self):
 162         dict = self.item().to_dict()
 163         dict['user'] = self.user
 164         return dict
 165
 166     def validate(self):
 167         ret = self.item().validate()
 168         if ret and self.item().version > self.version:
 169             self.version = self.item().version
 170             self.lib.save(self.id)
 171         return ret
 172
 173     def prepare(self):
 174         ret = self.item().prepare()
 175         if ret and self.item().version > self.version:
 176             self.version = self.item().version
 177             self.lib.save(self.id)
 178         return ret
 179
 180     def async_prepare(self):
 181         th = threading.Thread(
 182             target=self.prepare, name="Prepare-" + self.id[:7])
 183         self.log.info(
 184             "%s: start preparing item in thread: " % self.item().type + self.format_debug_string())
 185         th.daemon = True
 186         th.start()
 187         return th
 188
 189     def uri(self):
 190         return self.item().uri()
 191
 192     def add_tags(self, tags):
 193         self.item().add_tags(tags)
 194         if self.item().version > self.version:
 195             self.version = self.item().version
 196             self.lib.save(self.id)
 197
 198     def remove_tags(self, tags):
 199         self.item().remove_tags(tags)
 200         if self.item().version > self.version:
 201             self.version = self.item().version
 202             self.lib.save(self.id)
 203
 204     def clear_tags(self):
 205         self.item().clear_tags()
 206         if self.item().version > self.version:
 207             self.version = self.item().version
 208             self.lib.save(self.id)
 209
 210     def is_ready(self):
 211         return self.item().is_ready()
 212
 213     def is_failed(self):
 214         return self.item().is_failed()
 215
 216     def format_current_playing(self):
 217         return self.item().format_current_playing(self.user)
 218
 219     def format_song_string(self):
 220         return self.item().format_song_string(self.user)
 221
 222     def format_short_string(self):
 223         return self.item().format_short_string()
 224
 225     def format_debug_string(self):
 226         return self.item().format_debug_string()
 227
 228     def display_type(self):
 229         return self.item().display_type()
 230
 231
 232 # Remember!!! Get wrapper functions will automatically add items into the cache!
 233 def get_cached_wrapper_from_scrap(bot, **kwargs):
 234     item = var.cache.get_item(bot, **kwargs)
 235     if 'user' not in kwargs:
 236         raise KeyError("Which user added this song?")
 237     return CachedItemWrapper(var.cache, item.id, kwargs['type'], kwargs['user'])
 238
 239
 240 def get_cached_wrapper_from_dict(bot, dict_from_db, user):
 241     item = dict_to_item(bot, dict_from_db)
 242     var.cache[dict_from_db['id']] = item
 243     return CachedItemWrapper(var.cache, item.id, item.type, user)
 244
 245
 246 def get_cached_wrapper_by_id(bot, id, user):
 247     item = var.cache.get_item_by_id(bot, id)
 248     if item:
 249         return CachedItemWrapper(var.cache, item.id, item.type, user)
 250     else:
 251         return None
 252
 253
 254 def get_cached_wrappers_by_tags(bot, tags, user):
 255     items = var.cache.get_items_by_tags(bot, tags)
 256     ret = []
 257     for item in items:
 258         ret.append(CachedItemWrapper(var.cache, item.id, item.type, user))
 259     return ret