"""MySQL implementation of the database class.""" import queue import mysql.connector from mysql.connector.errors import DatabaseError from dejavu.base_classes.common_database import CommonDatabase from dejavu.config.settings import (FIELD_FILE_SHA1, FIELD_FINGERPRINTED, FIELD_EXTENSION, FIELD_HASH, FIELD_OFFSET, FIELD_SONG_ID, FIELD_SONGNAME, FIELD_TOTAL_HASHES, TOPN, FINGERPRINTS_TABLENAME, SONGS_TABLENAME) class MySQLDatabase(CommonDatabase): type = "mysql" # CREATES CREATE_SONGS_TABLE = f""" CREATE TABLE IF NOT EXISTS `{SONGS_TABLENAME}` ( `{FIELD_SONG_ID}` MEDIUMINT UNSIGNED NOT NULL AUTO_INCREMENT , `{FIELD_SONGNAME}` VARCHAR(250) NOT NULL , `{FIELD_EXTENSION}` VARCHAR(16) NOT NULL , `{FIELD_FINGERPRINTED}` TINYINT DEFAULT 0 , `{FIELD_FILE_SHA1}` BINARY(20) NOT NULL , `{FIELD_TOTAL_HASHES}` INT NOT NULL DEFAULT 0 , `date_created` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP , `date_modified` DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP , CONSTRAINT `pk_{SONGS_TABLENAME}_{FIELD_SONG_ID}` PRIMARY KEY (`{FIELD_SONG_ID}`) , CONSTRAINT `uq_{SONGS_TABLENAME}_{FIELD_SONG_ID}` UNIQUE KEY (`{FIELD_SONG_ID}`) , CONSTRAINT `uq_{SONGS_TABLENAME}-{FIELD_SONGNAME}` UNIQUE KEY (`{FIELD_SONGNAME}`) ) ENGINE=INNODB; """ CREATE_FINGERPRINTS_TABLE = f""" CREATE TABLE IF NOT EXISTS `{FINGERPRINTS_TABLENAME}` ( `{FIELD_HASH}` BINARY(10) NOT NULL , `{FIELD_SONG_ID}` MEDIUMINT UNSIGNED NOT NULL , `{FIELD_OFFSET}` INT UNSIGNED NOT NULL , INDEX `ix_{FINGERPRINTS_TABLENAME}_{FIELD_HASH}` (`{FIELD_HASH}`) , CONSTRAINT `uq_{FINGERPRINTS_TABLENAME}_{FIELD_SONG_ID}_{FIELD_OFFSET}_{FIELD_HASH}` UNIQUE KEY (`{FIELD_SONG_ID}`, `{FIELD_OFFSET}`, `{FIELD_HASH}`) , CONSTRAINT `fk_{FINGERPRINTS_TABLENAME}_{FIELD_SONG_ID}` FOREIGN KEY (`{FIELD_SONG_ID}`) REFERENCES `{SONGS_TABLENAME}`(`{FIELD_SONG_ID}`) ON DELETE CASCADE ) ENGINE=INNODB; """ # INSERTS (IGNORES DUPLICATES) INSERT_FINGERPRINT = f""" INSERT IGNORE INTO `{FINGERPRINTS_TABLENAME}` ( `{FIELD_SONG_ID}` , `{FIELD_HASH}` , `{FIELD_OFFSET}`) VALUES (%s, UNHEX(%s), %s); """ INSERT_SONG = f""" INSERT INTO `{SONGS_TABLENAME}` (`{FIELD_SONGNAME}`,`{FIELD_FILE_SHA1}`,`{FIELD_TOTAL_HASHES}`,`{FIELD_EXTENSION}`) VALUES (%s, UNHEX(%s), %s, %s); """ # SELECTS SELECT = f""" SELECT `{FIELD_SONG_ID}`, `{FIELD_OFFSET}` FROM `{FINGERPRINTS_TABLENAME}` WHERE `{FIELD_HASH}` = UNHEX(%s); """ SELECT_FINGERPRINTS = f""" SELECT HEX(`{FIELD_HASH}`), `{FIELD_SONG_ID}`, COUNT(`{FIELD_SONG_ID}`) as cid, GROUP_CONCAT(`{FIELD_OFFSET}`) as offsets FROM `{FINGERPRINTS_TABLENAME}` WHERE `{FIELD_HASH}` IN (%s) GROUP BY (`{FIELD_SONG_ID}`) ORDER BY cid DESC LIMIT {TOPN}; """ SELECT_FINGERPRINTS_EXCLUDE_ID = f""" SELECT HEX(`{FIELD_HASH}`), `{FIELD_SONG_ID}`, COUNT(`{FIELD_SONG_ID}`) as cid, GROUP_CONCAT(`{FIELD_OFFSET}`) as offsets FROM `{FINGERPRINTS_TABLENAME}` WHERE `{FIELD_HASH}` IN (%s) AND `{FIELD_SONG_ID}` != %s GROUP BY (`{FIELD_SONG_ID}`) ORDER BY cid DESC LIMIT {TOPN}; """ SELECT_ALL = f"SELECT `{FIELD_SONG_ID}`, `{FIELD_OFFSET}` FROM `{FINGERPRINTS_TABLENAME}`;" SELECT_SONG = f""" SELECT `{FIELD_SONGNAME}`, HEX(`{FIELD_FILE_SHA1}`) AS `{FIELD_FILE_SHA1}`, `{FIELD_TOTAL_HASHES}`, `{FIELD_EXTENSION}` FROM `{SONGS_TABLENAME}` WHERE `{FIELD_SONG_ID}` = %s; """ SELECT_SONGS_BY_ID = f""" SELECT `{FIELD_SONG_ID}`, `{FIELD_SONGNAME}`, HEX(`{FIELD_FILE_SHA1}`) AS `{FIELD_FILE_SHA1}`, `{FIELD_TOTAL_HASHES}`, `{FIELD_EXTENSION}` FROM `{SONGS_TABLENAME}` WHERE `{FIELD_SONG_ID}` IN (%s); """ SELECT_SONG_ID_BY_NAME = f""" SELECT `{FIELD_SONG_ID}` FROM `{SONGS_TABLENAME}` WHERE `{FIELD_SONGNAME}` = %s; """ SELECT_NUM_FINGERPRINTS = f"SELECT COUNT(*) AS n FROM `{FINGERPRINTS_TABLENAME}`;" SELECT_UNIQUE_SONG_IDS = f""" SELECT COUNT(`{FIELD_SONG_ID}`) AS n FROM `{SONGS_TABLENAME}` WHERE `{FIELD_FINGERPRINTED}` = 1; """ SELECT_SONGS = f""" SELECT `{FIELD_SONGNAME}` , `{FIELD_EXTENSION}` FROM `{SONGS_TABLENAME}`; """ # DROPS DROP_FINGERPRINTS = f"DROP TABLE IF EXISTS `{FINGERPRINTS_TABLENAME}`;" DROP_SONGS = f"DROP TABLE IF EXISTS `{SONGS_TABLENAME}`;" # UPDATE UPDATE_SONG_FINGERPRINTED = f""" UPDATE `{SONGS_TABLENAME}` SET `{FIELD_FINGERPRINTED}` = 1 WHERE `{FIELD_SONG_ID}` = %s; """ # DELETES DELETE_UNFINGERPRINTED = f""" DELETE FROM `{SONGS_TABLENAME}` WHERE `{FIELD_FINGERPRINTED}` = 0; """ DELETE_SONGS = f""" DELETE FROM `{SONGS_TABLENAME}` WHERE `{FIELD_SONG_ID}` IN (%s); """ # IN IN_MATCH = f"UNHEX(%s)" def __init__(self, **options): super().__init__() self.cursor = cursor_factory(**options) self._options = options def after_fork(self) -> None: # Clear the cursor cache, we don't want any stale connections from # the previous process. Cursor.clear_cache() def insert_song(self, song_name: str, file_hash: str, total_hashes: int, extension: str) -> int: """ Inserts a song name into the database, returns the new identifier of the song. :param song_name: The name of the audio file. :param file_hash: Hash from the fingerprinted file. :param total_hashes: amount of hashes to be inserted on fingerprint table. :param extension: file extension :return: the inserted id. """ with self.cursor() as cur: cur.execute(self.INSERT_SONG, (song_name, file_hash, total_hashes, extension)) return cur.lastrowid def get_cursor(self): return cursor_factory(**self._options) def __getstate__(self): return self._options, def __setstate__(self, state): self._options, = state self.cursor = cursor_factory(**self._options) def cursor_factory(**factory_options): def cursor(**options): options.update(factory_options) return Cursor(**options) return cursor class Cursor(object): """ Establishes a connection to the database and returns an open cursor. # Use as context manager with Cursor() as cur: cur.execute(query) ... """ def __init__(self, dictionary=False, **options): super().__init__() self._cache = queue.Queue(maxsize=50) try: conn = self._cache.get_nowait() # Ping the connection before using it from the cache. conn.ping(True) except queue.Empty: conn = mysql.connector.connect(**options) self.conn = conn self.dictionary = dictionary @classmethod def clear_cache(cls): cls._cache = queue.Queue(maxsize=50) def __enter__(self): self.cursor = self.conn.cursor(dictionary=self.dictionary) return self.cursor def __exit__(self, extype, exvalue, traceback): # if we had a MySQL related error we try to rollback the cursor. if extype is DatabaseError: self.cursor.rollback() self.cursor.close() self.conn.commit() # Put it back on the queue try: self._cache.put_nowait(self.conn) except queue.Full: self.conn.close()