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