mirror of
https://github.com/nathom/streamrip.git
synced 2024-09-19 19:28:46 -04:00
77 lines
2.1 KiB
Python
77 lines
2.1 KiB
Python
"""Wrapper over a database that stores item IDs."""
|
|
|
|
import logging
|
|
import os
|
|
import sqlite3
|
|
from typing import Union
|
|
|
|
logger = logging.getLogger("streamrip")
|
|
|
|
|
|
class MusicDB:
|
|
"""Simple interface for the downloaded track database."""
|
|
|
|
def __init__(self, db_path: Union[str, os.PathLike], empty=False):
|
|
"""Create a MusicDB object.
|
|
|
|
:param db_path: filepath of the database
|
|
:type db_path: Union[str, os.PathLike]
|
|
"""
|
|
if empty:
|
|
self.path = None
|
|
|
|
self.path = db_path
|
|
if not os.path.exists(self.path):
|
|
self.create()
|
|
|
|
def create(self):
|
|
"""Create a database at `self.path`."""
|
|
with sqlite3.connect(self.path) as conn:
|
|
try:
|
|
conn.execute("CREATE TABLE downloads (id TEXT UNIQUE NOT NULL);")
|
|
logger.debug("Download-IDs database created: %s", self.path)
|
|
except sqlite3.OperationalError:
|
|
pass
|
|
|
|
return self.path
|
|
|
|
def __contains__(self, item_id: Union[str, int]) -> bool:
|
|
"""Check whether the database contains an id.
|
|
|
|
:param item_id: the id to check
|
|
:type item_id: str
|
|
:rtype: bool
|
|
"""
|
|
if self.path is None:
|
|
return False
|
|
|
|
logger.debug("Checking database for ID %s", item_id)
|
|
with sqlite3.connect(self.path) as conn:
|
|
return (
|
|
conn.execute(
|
|
"SELECT id FROM downloads where id=?", (item_id,)
|
|
).fetchone()
|
|
is not None
|
|
)
|
|
|
|
def add(self, item_id: str):
|
|
"""Add an id to the database.
|
|
|
|
:param item_id:
|
|
:type item_id: str
|
|
"""
|
|
logger.debug("Adding ID %s", item_id)
|
|
|
|
if self.path is None:
|
|
return
|
|
|
|
with sqlite3.connect(self.path) as conn:
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO downloads (id) VALUES (?)",
|
|
(item_id,),
|
|
)
|
|
conn.commit()
|
|
except sqlite3.Error as err:
|
|
if "UNIQUE" not in str(err):
|
|
raise
|