Add support for creating indices on arbitrary database tables

beets/dbcore/__init__.py

@@ -16,7 +16,7 @@
 Library.
 """
 
-from .db import Database, Model, Results
+from .db import Database, Index, Model, Results
 from .query import (
     AndQuery,
     FieldQuery,
@@ -43,6 +43,7 @@
     "Query",
     "Results",
     "Type",
+    "Index",
     "parse_sorted_query",
     "query_from_strings",
     "sort_from_strings",

beets/dbcore/db.py

@@ -26,7 +26,7 @@
 from collections import defaultdict
 from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence
 from sqlite3 import Connection
-from typing import TYPE_CHECKING, Any, AnyStr, Callable, Generic
+from typing import TYPE_CHECKING, Any, AnyStr, Callable, Generic, NamedTuple
 
 from typing_extensions import TypeVar  # default value support
 from unidecode import unidecode
@@ -287,6 +287,11 @@ class Model(ABC, Generic[D]):
     terms.
     """
 
+    _indices: Sequence[Index] = ()
+    """A sequence of `Index` objects that describe the indices to be
+    created for this table.
+    """
+
     @cached_classproperty
     def _types(cls) -> dict[str, types.Type]:
         """Optional types for non-fixed (flexible and computed) fields."""
@@ -1033,6 +1038,7 @@ def __init__(self, path, timeout: float = 5.0):
         for model_cls in self._models:
             self._make_table(model_cls._table, model_cls._fields)
             self._make_attribute_table(model_cls._flex_table)
+            self._migrate_indices(model_cls._table, model_cls._indices)
 
     # Primitive access control: connections and transactions.
 
@@ -1207,6 +1213,25 @@ def _make_attribute_table(self, flex_table: str):
                 """.format(flex_table)
             )
 
+    def _migrate_indices(
+        self,
+        table: str,
+        indices: Sequence[Index],
+    ):
+        """Create or replace indices for the given table.
+
+        If the indices already exists and are up to date (i.e., the
+        index name and columns match), nothing is done. Otherwise, the
+        indices are created or replaced.
+        """
+        with self.transaction() as tx:
+            current = {
+                Index.from_db(tx, r[1])
+                for r in tx.query(f"PRAGMA index_list({table})")
+            }
+            for index in set(indices) - current:
+                index.recreate(tx, table)
+
     # Querying.
 
     def _fetch(
@@ -1276,3 +1301,38 @@ def _get(
         exist.
         """
         return self._fetch(model_cls, MatchQuery("id", id)).get()
+
+
+class Index(NamedTuple):
+    """A helper class to represent the index
+    information in the database schema.
+    """
+
+    name: str
+    columns: tuple[str, ...]
+
+    def recreate(self, tx: Transaction, table: str) -> None:
+        """Recreate the index in the database.
+
+        This is useful when the index has been changed and needs to be
+        updated.
+        """
+        tx.script(f"""
+            DROP INDEX IF EXISTS {self.name};
+            CREATE INDEX {self.name} ON {table} ({", ".join(self.columns)})
+        """)
+
+    @classmethod
+    def from_db(cls, tx: Transaction, name: str) -> Index:
+        """Create an Index object from the database if it exists.
+
+        The name has to exists in the database! Otherwise, an
+        Error will be raised.
+        """
+        rows = tx.query(f"PRAGMA index_info({name})")
+        columns = tuple(row[2] for row in rows)
+        return cls(name, columns)
+
+    def __hash__(self) -> int:
+        """Unique hash for the index based on its name and columns."""
+        return hash((self.name, tuple(self.columns)))

beets/library/models.py

@@ -716,6 +716,7 @@ class Item(LibModel):
         "mtime": types.DATE,
         "added": types.DATE,
     }
+    _indices = (dbcore.Index("idx_item_album_id", ("album_id",)),)
 
     _search_fields = (
         "artist",

docs/changelog.rst

@@ -52,6 +52,9 @@ Bug fixes:
   the config option ``spotify.search_query_ascii: yes``. :bug:`5699`
 - :doc:`plugins/discogs`: Beets will no longer crash if a release has been
   deleted, and returns a 404.
+- Beets now creates an index for the ``album_id`` field in the ``items`` table.
+  This significantly speeds up queries that filter items by their album.
+  :bug:`5809`
 
 For packagers:
 

test/test_dbcore.py

@@ -23,6 +23,7 @@
 import pytest
 
 from beets import dbcore
+from beets.dbcore.db import Index
 from beets.library import LibModel
 from beets.test import _common
 from beets.util import cached_classproperty
@@ -58,6 +59,7 @@ class ModelFixture1(LibModel):
     _sorts = {
         "some_sort": SortFixture,
     }
+    _indices = (Index("field_one_index", ("field_one",)),)
 
     @cached_classproperty
     def _types(cls):
@@ -133,6 +135,7 @@ class AnotherModelFixture(ModelFixture1):
         "id": dbcore.types.PRIMARY_ID,
         "foo": dbcore.types.INTEGER,
     }
+    _indices = (Index("another_foo_index", ("foo",)),)
 
 
 class ModelFixture5(ModelFixture1):
@@ -784,3 +787,48 @@ def test_no_results(self):
             self.db._fetch(ModelFixture1, dbcore.query.FalseQuery()).get()
             is None
         )
+
+
+class TestIndex:
+    @pytest.fixture(autouse=True)
+    def db(self):
+        """Set up an in-memory SQLite database."""
+        db = DatabaseFixture1(":memory:")
+        yield db
+        db._connection().close()
+
+    @pytest.fixture
+    def sample_index(self):
+        """Fixture for a sample Index object."""
+        return Index(name="sample_index", columns=("field_one",))
+
+    def test_index_creation(self, db, sample_index):
+        """Test creating an index and checking its existence."""
+        with db.transaction() as tx:
+            sample_index.recreate(tx, "test")
+            indexes = (
+                db._connection().execute("PRAGMA index_list(test)").fetchall()
+            )
+        assert any(sample_index.name in index for index in indexes)
+
+    def test_from_db(self, db, sample_index):
+        """Test retrieving an index from the database."""
+        with db.transaction() as tx:
+            sample_index.recreate(tx, "test")
+            retrieved = Index.from_db(tx, sample_index.name)
+            assert retrieved == sample_index
+
+    def test_index_hashing_and_set_behavior(self, sample_index):
+        """Test the hashing and set behavior of the Index class."""
+        index_set = {sample_index}
+        similar_index = Index(name="sample_index", columns=("field_one",))
+
+        assert similar_index in index_set  # Should recognize similar attributes
+
+        different_index = Index(name="other_index", columns=("other_field",))
+        index_set.add(different_index)
+
+        assert len(index_set) == 2  # Should recognize distinct index
+
+        index_set.discard(sample_index)
+        assert similar_index not in index_set  # Should remove similar index