pyrekordbox 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

pyrekordbox/db6/database.py

@@ -9,9 +9,10 @@ from uuid import uuid4
 from pathlib import Path
 from typing import Optional
 from sqlalchemy import create_engine, or_, event, MetaData
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.orm import Session
 from sqlalchemy.exc import NoResultFound
-from packaging import version
+from sqlalchemy.sql.sqltypes import DateTime, String
+import packaging.version
 from ..utils import get_rekordbox_pid
 from ..config import get_config
 from ..anlz import get_anlz_paths, read_anlz_files
@@ -21,14 +22,11 @@ from .tables import DjmdContent
 from . import tables
 
 try:
-    from pysqlcipher3 import dbapi2 as sqlite3  # noqa
+    from sqlcipher3 import dbapi2 as sqlite3  # noqa
 except ImportError:
-    try:
-        from sqlcipher3 import dbapi2 as sqlite3  # noqa
-    except ImportError:
-        import sqlite3
+    import sqlite3
 
-MAX_VERSION = version.parse("6.6.5")
+MAX_VERSION = packaging.version.parse("6.6.5")
 
 logger = logging.getLogger(__name__)
 
@@ -87,12 +85,12 @@ def open_rekordbox_database(path=None, key="", unlock=True, sql_driver=None):
 
     To use the ``pysqlcipher3`` package as SQLite driver, either import it as
 
-    >>> from pysqlcipher3 import dbapi2 as sqlite3  # noqa
+    >>> from sqlcipher3 import dbapi2 as sqlite3  # noqa
     >>> db = open_rekordbox_database("path/to/master_copy.db")
 
     or supply the package as driver:
 
-    >>> from pysqlcipher3 import dbapi2
+    >>> from sqlcipher3 import dbapi2  # noqa
     >>> db = open_rekordbox_database("path/to/master_copy.db", sql_driver=dbapi2)
     """
     if not path:
@@ -111,7 +109,7 @@ def open_rekordbox_database(path=None, key="", unlock=True, sql_driver=None):
 
     if unlock:
         if not key:
-            ver = version.parse(rb6_config["version"])
+            ver = packaging.version.parse(rb6_config["version"])
             if ver >= MAX_VERSION:
                 raise IncompatibleVersionError(rb6_config["version"])
             try:
@@ -209,15 +207,15 @@ class Rekordbox6Database:
             raise FileNotFoundError(f"File '{path}' does not exist!")
         # Open database
         if unlock:
+            if "dp" in rb6_config:
+                key = rb6_config["dp"]
             if not key:
-                ver = version.parse(rb6_config["version"])
+                ver = packaging.version.parse(rb6_config["version"])
                 if ver >= MAX_VERSION:
                     raise IncompatibleVersionError(rb6_config["version"])
-                try:
-                    key = rb6_config["dp"]
-                except KeyError:
+                else:
                     raise ValueError("Could not unlock database: No key found")
-                logger.info("Key: %s", key)
+            logger.info("Key: %s", key)
             # Unlock database and create engine
             url = f"sqlite+pysqlcipher://:{key}@/{path}?"
             engine = create_engine(url, module=sqlite3)
@@ -231,7 +229,6 @@ class Rekordbox6Database:
             raise FileNotFoundError(f"Database directory '{db_dir}' does not exist!")
 
         self.engine = engine
-        self._Session = sessionmaker(bind=self.engine)
         self.session: Optional[Session] = None
 
         self.registry = RekordboxAgentRegistry(self)
@@ -272,7 +269,7 @@ class Rekordbox6Database:
         >>> db.open()
         """
         if self.session is None:
-            self.session = self._Session()
+            self.session = Session(bind=self.engine)
             self.registry.clear_buffer()
 
     def close(self):
@@ -928,12 +925,12 @@ class Rekordbox6Database:
             .order_by(tables.DjmdSongPlaylist.TrackNo)
         )
         moved = list()
-        self.registry.disable_tracking()
-        for song in query:
-            song.TrackNo -= 1
-            song.updated_at = now
-            moved.append(song)
-        self.registry.enable_tracking()
+        with self.registry.disabled():
+            for song in query:
+                song.TrackNo -= 1
+                song.updated_at = now
+                moved.append(song)
+
         if moved:
             self.registry.on_move(moved)
 
@@ -1095,9 +1092,8 @@ class Rekordbox6Database:
             )
             for pl in query:
                 pl.Seq += 1
-                self.registry.disable_tracking()
-                pl.updated_at = now
-                self.registry.enable_tracking()
+                with self.registry.disabled():
+                    pl.updated_at = now
 
         # Add new playlist to database
         # First create with name 'New playlist'
@@ -1393,10 +1389,9 @@ class Rekordbox6Database:
             # Set seq number and update time *before* other playlists to ensure
             # right USN increment order
             playlist.ParentID = parent_id
-            self.registry.disable_tracking()
-            playlist.Seq = seq
-            playlist.updated_at = now
-            self.registry.enable_tracking()
+            with self.registry.disabled():
+                playlist.Seq = seq
+                playlist.updated_at = now
 
             if not insert_at_end:
                 # Update seq numbers higher than the new seq number in *new* parent
@@ -1405,9 +1400,8 @@ class Rekordbox6Database:
                     # Update time of other playlists are left unchanged
                     pl.Seq += 1
                     # Each move counts as one USN increment, so disable for update time
-                    self.registry.disable_tracking()
-                    pl.updated_at = now
-                    self.registry.enable_tracking()
+                    with self.registry.disabled():
+                        pl.updated_at = now
 
             # Update seq numbers higher than the old seq number in *old* parent
             # USN is not updated here
@@ -1466,17 +1460,15 @@ class Rekordbox6Database:
             # right USN increment order
             playlist.Seq = seq
             # Each move counts as one USN increment, so disable for update time
-            self.registry.disable_tracking()
-            playlist.updated_at = now
-            self.registry.enable_tracking()
+            with self.registry.disabled():
+                playlist.updated_at = now
 
             # Set seq number and update time for playlists between old_seq and seq
             for pl in other_playlists:
                 pl.Seq += delta_seq
                 # Each move counts as one USN increment, so disable for update time
-                self.registry.disable_tracking()
-                pl.updated_at = now
-                self.registry.enable_tracking()
+                with self.registry.disabled():
+                    pl.updated_at = now
 
     def rename_playlist(self, playlist, name):
         """Renames a playlist or playlist folder.
@@ -1512,9 +1504,241 @@ class Rekordbox6Database:
         # Update name of playlist
         playlist.Name = name
         # Update update time: USN not incremented
-        self.registry.disable_tracking()
-        playlist.updated_at = now
-        self.registry.enable_tracking()
+        with self.registry.disabled():
+            playlist.updated_at = now
+
+    def add_album(
+        self, name, artist=None, image_path=None, compilation=None, search_str=None
+    ):
+        """Adds a new album to the database.
+
+        Parameters
+        ----------
+        name : str
+            The name of the album. Must be a unique name (case-sensitive).
+            If an album with the same name already exists in the database,
+            use the `ID` of the existing album instead.
+        artist : str or int or DjmdArtist, optional
+            The artist of the album. Can either be a :class:`DjmdArtist` object
+            or an artist ID.
+        image_path : str, optional
+            The path to the album cover image.
+        compilation : bool, optional
+            Whether the album is a compilation album. If not given, the
+            default value of `False` is used.
+        search_str : str, optional
+            The search string of the album.
+
+        Returns
+        -------
+        album : DjmdAlbum
+            The newly created album.
+
+        Raises
+        ------
+        ValueError : If an album with the same name already exists in the database.
+
+        Examples
+        --------
+        Add a new album to the database:
+
+        >>> db = Rekordbox6Database()
+        >>> db.add_album(name="Album 1")
+        <DjmdAlbum(148754249  Name=Album 1)>
+
+        Add a new album to the database with an album artist:
+
+        >>> artist = db.get_artist(Name="Artist 1").one()  # noqa
+        >>> db.add_album(name="Album 2", artist=artist)
+        <DjmdAlbum(148754249  Name=Album 2)>
+
+        For setting the album of a track, the usual procedure is to first
+        check if an entry with the same album name already exists in the database,
+        and if not, add a new album:
+
+        >>> name = "Album name"
+        >>> content = db.get_content().one()
+        >>> album = db.get_album(Name=name).one_or_none()
+        >>> if album is None:
+        ...     album = db.add_album(name=name)
+        >>> content.AlbumID = album.ID
+        """
+        # Check if album already exists
+        query = self.query(tables.DjmdAlbum).filter_by(Name=name)
+        if query.count() > 0:
+            raise ValueError(f"Album '{name}' already exists in database")
+
+        # Get artist ID
+        if artist is not None:
+            if isinstance(artist, (int, str)):
+                artist = self.get_artist(ID=artist)
+            artist = artist.ID
+
+        album = tables.DjmdAlbum.create(
+            ID=self.generate_unused_id(tables.DjmdAlbum),
+            Name=name,
+            AlbumArtistID=artist,
+            ImagePath=image_path,
+            Compilation=compilation,
+            SearchStr=search_str,
+        )
+        self.add(album)
+        self.flush()
+        return album
+
+    def add_artist(self, name, search_str=None):
+        """Adds a new artist to the database.
+
+        Parameters
+        ----------
+        name : str
+            The name of the artist. Must be a unique name (case-sensitive).
+            If an artist with the same name already exists in the database,
+            use the `ID` of the existing artist instead.
+        search_str : str, optional
+            The search string of the artist.
+
+        Returns
+        -------
+        artist : DjmdArtist
+            The newly created artist.
+
+        Raises
+        ------
+        ValueError : If an artist with the same name already exists in the database.
+
+        Examples
+        --------
+        Add a new artist to the database:
+
+        >>> db = Rekordbox6Database()
+        >>> db.add_artist(name="Artist 1")
+        <DjmdArtist(123456789, Name='Artist 1')>
+
+        Add a new artist to the database with a custom search string:
+
+        >>> db.add_artist(name="Artist 2", search_str="artist 2")
+        <DjmdArtist(123456789, Name='Artist 2')>
+
+        For setting the artist of a track, the usual procedure is to first
+        check if an entry with the same artist name already exists in the database,
+        and if not, add a new artist:
+
+        >>> name = "Artist name"
+        >>> content = db.get_content().one()
+        >>> artist = db.get_artist(Name=name).one_or_none()
+        >>> if artist is None:
+        ...     artist = db.add_artist(name=name)
+        >>> content.ArtistID = artist.ID
+        """
+        # Check if artist already exists
+        query = self.query(tables.DjmdArtist).filter_by(Name=name)
+        if query.count() > 0:
+            raise ValueError(f"Artist '{name}' already exists in database")
+
+        id_ = self.generate_unused_id(tables.DjmdArtist)
+        artist = tables.DjmdArtist.create(ID=id_, Name=name, SearchStr=search_str)
+        self.add(artist)
+        self.flush()
+        return artist
+
+    def add_genre(self, name):
+        """Adds a new genre to the database.
+
+        Parameters
+        ----------
+        name : str
+            The name of the genre. Must be a unique name (case-sensitive).
+            If a genre with the same name already exists in the database,
+            use the `ID` of the existing genre instead.
+
+        Returns
+        -------
+        genre : DjmdGenre
+            The newly created genre.
+
+        Raises
+        ------
+        ValueError : If a genre with the same name already exists in the database.
+
+        Examples
+        --------
+        Add a new genre to the database:
+
+        >>> db = Rekordbox6Database()
+        >>> db.add_genre(name="Genre 1")
+        <DjmdGenre(123456789 Name=Genre 1)>
+
+        For setting the genre of a track, the usual procedure is to first
+        check if an entry with the same genre name already exists in the database,
+        and if not, add a new genre:
+
+        >>> name = "Genre name"
+        >>> content = db.get_content().one()
+        >>> genre = db.get_genre(Name=name).one_or_none()
+        >>> if genre is None:
+        ...     genre = db.add_genre(name=name)
+        >>> content.GenreID = genre.ID
+        """
+        # Check if genre already exists
+        query = self.query(tables.DjmdGenre).filter_by(Name=name)
+        if query.count() > 0:
+            raise ValueError(f"Genre '{name}' already exists in database")
+
+        id_ = self.generate_unused_id(tables.DjmdGenre)
+        genre = tables.DjmdGenre.create(ID=id_, Name=name)
+        self.add(genre)
+        self.flush()
+        return genre
+
+    def add_label(self, name):
+        """Adds a new label to the database.
+
+        Parameters
+        ----------
+        name : str
+            The name of the label. Must be a unique name (case-sensitive).
+            If a label with the same name already exists in the database,
+            use the `ID` of the existing label instead.
+
+        Returns
+        -------
+        label : DjmdLabel
+            The newly created label.
+
+        Raises
+        ------
+        ValueError : If a label with the same name already exists in the database.
+
+        Examples
+        --------
+        Add a new label to the database:
+
+        >>> db = Rekordbox6Database()
+        >>> db.add_label(name="Label 1")
+        <DjmdLabel(123456789 Name=Label 1)>
+
+        For setting the label of a track, the usual procedure is to first
+        check if an entry with the same label name already exists in the database,
+        and if not, add a new label:
+
+        >>> name = "Label name"
+        >>> content = db.get_content().one()
+        >>> label = db.get_label(Name=name).one_or_none()
+        >>> if label is None:
+        ...     label = db.add_label(name=name)
+        >>> content.LabelID = label.ID
+        """
+        # Check if label already exists
+        query = self.query(tables.DjmdLabel).filter_by(Name=name)
+        if query.count() > 0:
+            raise ValueError(f"Label '{name}' already exists in database")
+
+        id_ = self.generate_unused_id(tables.DjmdLabel)
+        label = tables.DjmdLabel.create(ID=id_, Name=name)
+        self.add(label)
+        self.flush()
+        return label
 
     # ----------------------------------------------------------------------------------
 
@@ -1761,46 +1985,51 @@ class Rekordbox6Database:
             json.dump(data, fp, indent=indent, sort_keys=sort_keys, default=json_serial)
 
     def copy_unlocked(self, output_file):
-        src_metadata = MetaData(bind=self.engine)
+        src_engine = self.engine
+        src_metadata = MetaData()
         exclude_tables = ("sqlite_master", "sqlite_sequence", "sqlite_temp_master")
 
         dst_engine = create_engine(f"sqlite:///{output_file}")
-        dst_metadata = MetaData(bind=dst_engine)
+        dst_metadata = MetaData()
 
         @event.listens_for(src_metadata, "column_reflect")
         def genericize_datatypes(inspector, tablename, column_dict):
-            column_dict["type"] = column_dict["type"].as_generic(allow_nulltype=True)
-
-        dst_engine.connect()
-        dst_metadata.reflect()
-
+            type_ = column_dict["type"].as_generic(allow_nulltype=True)
+            if isinstance(type_, DateTime):
+                type_ = String
+            column_dict["type"] = type_
+
+        src_conn = src_engine.connect()
+        dst_conn = dst_engine.connect()
+        dst_metadata.reflect(bind=dst_engine)
         # drop all tables in target database
         for table in reversed(dst_metadata.sorted_tables):
             if table.name not in exclude_tables:
                 print("dropping table =", table.name)
-                table.drop()
-
+                table.drop(bind=dst_engine)
+        # Delete all data in target database
+        for table in reversed(dst_metadata.sorted_tables):
+            table.delete()
         dst_metadata.clear()
-        dst_metadata.reflect()
-        src_metadata.reflect()
-
+        dst_metadata.reflect(bind=dst_engine)
+        src_metadata.reflect(bind=src_engine)
         # create all tables in target database
         for table in src_metadata.sorted_tables:
             if table.name not in exclude_tables:
                 table.create(bind=dst_engine)
-
         # refresh metadata before you can copy data
         dst_metadata.clear()
-        dst_metadata.reflect()
-
+        dst_metadata.reflect(bind=dst_engine)
         # Copy all data from src to target
         print("Copying data...")
         string = "\rCopying table {name}: Inserting row {row}"
+        index = 0
         for table in dst_metadata.sorted_tables:
             src_table = src_metadata.tables[table.name]
             stmt = table.insert()
-            index = 0
-            for index, row in enumerate(src_table.select().execute()):
+            for index, row in enumerate(src_conn.execute(src_table.select())):
                 print(string.format(name=table.name, row=index), end="", flush=True)
-                stmt.execute(row._asdict())  # noqa
+                dst_conn.execute(stmt.values(row))
             print(f"\rCopying table {table.name}: Inserted {index} rows", flush=True)
+
+        dst_conn.commit()

pyrekordbox/db6/registry.py

@@ -3,6 +3,7 @@
 # Date:   2023-08-07
 
 import logging
+from contextlib import contextmanager
 
 logger = logging.getLogger(__name__)
 
@@ -103,6 +104,27 @@ class RekordboxAgentRegistry:
         """Disables the tracking of database changes."""
         cls.__enabled__ = False
 
+    @classmethod
+    @contextmanager
+    def disabled(cls):
+        """Context manager to temporarily disable the tracking of database changes.
+
+        Examples
+        --------
+        >>> registry = RekordboxAgentRegistry(db)  # noqa
+        >>> registry.__enabled__
+        True
+
+        >>> with registry.disabled():
+        ...     print(registry.__enabled__)
+        False
+        """
+        enabled = cls.__enabled__
+        cls.disable_tracking()
+        yield cls
+        if enabled:
+            cls.enable_tracking()
+
     def get_registries(self):
         """Returns all agent registries.
 
@@ -285,6 +307,7 @@ class RekordboxAgentRegistry:
         reg = self.db.get_agent_registry(registry_id="localUpdateCount")
         usn = reg.int_1
         self.disable_tracking()
+        self.db.flush()
         with self.db.session.no_autoflush:
             for instances, op, _, _ in self.__update_sequence__.copy():
                 usn += 1
@@ -298,5 +321,6 @@ class RekordboxAgentRegistry:
             reg.int_1 = usn
 
         self.clear_buffer()
+        self.db.flush()
         self.enable_tracking()
         return usn