Enable querying related flexible attributes

Author: snejus

beets/dbcore/db.py

@@ -743,33 +743,6 @@ def set_parse(self, key, string: str):
         """Set the object's key to a value represented by a string."""
         self[key] = self._parse(key, string)
 
-    # Convenient queries.
-
-    @classmethod
-    def field_query(
-        cls,
-        field,
-        pattern,
-        query_cls: Type[FieldQuery] = MatchQuery,
-    ) -> FieldQuery:
-        """Get a `FieldQuery` for this model."""
-        return query_cls(field, pattern, field in cls._fields)
-
-    @classmethod
-    def all_fields_query(
-        cls: Type["Model"],
-        pats: Mapping,
-        query_cls: Type[FieldQuery] = MatchQuery,
-    ):
-        """Get a query that matches many fields with different patterns.
-
-        `pats` should be a mapping from field names to patterns. The
-        resulting query is a conjunction ("and") of per-field queries
-        for all of these field/pattern pairs.
-        """
-        subqueries = [cls.field_query(k, v, query_cls) for k, v in pats.items()]
-        return AndQuery(subqueries)
-
 
 # Database controller and supporting interfaces.
 
@@ -1246,7 +1219,7 @@ def _fetch(
 
         relation_fields = query.model_fields - set(model_cls._fields)
         _from = model_cls.table_with_flex_attrs
-        if relation_fields:
+        if relation_fields or query.flex_fields:
             _from += f" {model_cls.relation_join}"
 
         table = model_cls._table_name

beets/dbcore/query.py

@@ -92,6 +92,11 @@ def model_fields(self) -> Set[str]:
         """Return query fields that are (any) model attributes."""
         return {f for f, fast in self.fields_info if fast}
 
+    @property
+    def flex_fields(self) -> Set[str]:
+        """Return query fields that are (any) model attributes."""
+        return {f for f, fast in self.fields_info if not fast}
+
     def clause(self) -> Tuple[Optional[str], Sequence[Any]]:
         """Generate an SQLite expression implementing the query.
 
@@ -149,7 +154,10 @@ def fields_info(self) -> Set[Tuple[str, bool]]:
     @property
     def col_name(self) -> str:
         if not self.fast:
-            return f'json_extract("flex_attrs [json_str]", "$.{self.field}")'
+            column = '"flex_attrs [json_str]"'
+            if self.table:
+                column = f"{self.table}.{column}"
+            return f'json_extract({column}, "$.{self.field}")'
 
         return f"{self.table}.{self.field}" if self.table else self.field
 

beets/dbcore/queryparse.py

@@ -16,12 +16,23 @@
 
 import itertools
 import re
-from typing import Collection, Dict, List, Optional, Sequence, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Collection,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+)
 
-from .. import library
 from . import Model, query
 from .query import Sort
 
+if TYPE_CHECKING:
+    from ..library import LibModel
+
 PARSE_QUERY_PART_REGEX = re.compile(
     # Non-capturing optional segment for the keyword.
     r"(-|\^)?"  # Negation prefixes.
@@ -105,7 +116,7 @@ def parse_query_part(
 
 
 def construct_query_part(
-    model_cls: Type[Model],
+    model_cls: Type["LibModel"],
     prefixes: Dict,
     query_part: str,
 ) -> query.Query:
@@ -153,17 +164,7 @@ def construct_query_part(
     # Field queries get constructed according to the name of the field
     # they are querying.
     else:
-        key = key.lower()
-        album_fields = library.Album._fields.keys()
-        item_fields = library.Item._fields.keys()
-        fast = key in album_fields | item_fields
-
-        if key in album_fields & item_fields:
-            # This field exists in both tables, so SQLite will encounter
-            # an OperationalError. Using an explicit table name resolves this.
-            key = f"{model_cls._table_name}.{key}"
-
-        out_query = query_class(key, pattern, fast)
+        out_query = model_cls.field_query(key.lower(), pattern, query_class)
 
     # Apply negation.
     if negate:

beets/library.py

@@ -22,6 +22,7 @@
 import sys
 import time
 import unicodedata
+from typing import Mapping, Set, Type
 
 from mediafile import MediaFile, UnreadableFileError
 
@@ -387,6 +388,14 @@ class LibModel(dbcore.Model):
     # Config key that specifies how an instance should be formatted.
     _format_config_key = None
 
+    @cached_classproperty
+    def all_model_db_fields(cls) -> Set[str]:
+        return Album._fields.keys() | Item._fields.keys()
+
+    @cached_classproperty
+    def shared_model_db_fields(cls) -> Set[str]:
+        return Album._fields.keys() & Item._fields.keys()
+
     def _template_funcs(self):
         funcs = DefaultTemplateFunctions(self, self._db).functions()
         funcs.update(plugins.template_funcs())
@@ -416,6 +425,46 @@ def __str__(self):
     def __bytes__(self):
         return self.__str__().encode("utf-8")
 
+    # Convenient queries.
+
+    @classmethod
+    def field_query(
+        cls, field: str, pattern: str, query_cls: Type[dbcore.FieldQuery]
+    ) -> dbcore.Query:
+        """Get a `FieldQuery` for this model."""
+        if field not in cls.all_model_db_fields:
+            # this is a flexible attribute. Since we do not know which entity
+            # it belongs, we check them both
+            tables = [Item._table_name, Album._table_name]
+            return dbcore.OrQuery(
+                [query_cls(f"{t}.{field}", pattern, False) for t in tables]
+            )
+
+        if field in cls.shared_model_db_fields:
+            # This field exists in both tables, so SQLite will encounter
+            # an OperationalError if we try to use it in a query.
+            # Using an explicit table name resolves this.
+            field = f"{cls._table_name}.{field}"
+
+        return query_cls(field, pattern, True)
+
+    @classmethod
+    def all_fields_query(
+        cls, pattern_by_field: Mapping[str, str]
+    ) -> dbcore.AndQuery:
+        """Get a query that matches many fields with different patterns.
+
+        `pattern_by_field` should be a mapping from field names to patterns.
+        The resulting query is a conjunction ("and") of per-field queries
+        for all of these field/pattern pairs.
+        """
+        return dbcore.AndQuery(
+            [
+                cls.field_query(f, p, dbcore.MatchQuery)
+                for f, p in pattern_by_field.items()
+            ]
+        )
+
 
 class FormattedItemMapping(dbcore.db.FormattedMapping):
     """Add lookup for album-level fields.
@@ -648,8 +697,7 @@ def relation_join(cls) -> str:
         We need to use a LEFT JOIN here, otherwise items that are not part of
         an album (e.g. singletons) would be left out.
         """
-        other_table = Album._table_name
-        return f"LEFT JOIN {other_table} ON {cls._table_name}.album_id = {other_table}.id"
+        return f"LEFT JOIN {Album.table_with_flex_attrs} ON {cls._table_name}.album_id = {Album._table_name}.id"
 
     @property
     def _cached_album(self):
@@ -1258,8 +1306,7 @@ def relation_join(cls) -> str:
         Here we can use INNER JOIN (which is more performant than LEFT JOIN),
         since we only want to see albums that
         """
-        other_table = Item._table_name
-        return f"INNER JOIN {other_table} ON {cls._table_name}.id = {other_table}.album_id"
+        return f"INNER JOIN {Item.table_with_flex_attrs} ON {cls._table_name}.id = {Item._table_name}.album_id"
 
     @classmethod
     def _getters(cls):
@@ -1949,9 +1996,10 @@ def _tmpl_unique(
             subqueries.extend(initial_subqueries)
         for key in keys:
             value = db_item.get(key, "")
-            # Use slow queries for flexible attributes.
-            fast = key in item_keys
-            subqueries.append(dbcore.MatchQuery(key, value, fast))
+            subqueries.append(
+                db_item.field_query(key, value, dbcore.MatchQuery)
+            )
+
         query = dbcore.AndQuery(subqueries)
         ambigous_items = (
             self.lib.items(query)

test/test_dbcore.py

@@ -21,6 +21,7 @@
 from tempfile import mkstemp
 
 from beets import dbcore
+from beets.library import LibModel
 from beets.test import _common
 
 # Fixture: concrete database and model classes. For migration tests, we
@@ -42,7 +43,7 @@ def match(self):
         return True
 
 
-class ModelFixture1(dbcore.Model):
+class ModelFixture1(LibModel):
     _table_name = "test"
     _flex_table = "testflex"
     _fields = {
@@ -590,15 +591,17 @@ def test_two_parts(self):
         self.assertIsInstance(q, dbcore.query.AndQuery)
         self.assertEqual(len(q.subqueries), 2)
         self.assertIsInstance(q.subqueries[0], dbcore.query.AnyFieldQuery)
-        self.assertIsInstance(q.subqueries[1], dbcore.query.SubstringQuery)
+        self.assertIsInstance(q.subqueries[1], dbcore.query.OrQuery)
+        self.assertIsInstance(q.subqueries[1][0], dbcore.query.SubstringQuery)
 
     def test_parse_fixed_type_query(self):
         q = self.qfs(["field_one:2..3"])
-        self.assertIsInstance(q.subqueries[0], dbcore.query.NumericQuery)
+        self.assertIsInstance(q.subqueries[0], dbcore.query.OrQuery)
+        self.assertIsInstance(q.subqueries[0][0], dbcore.query.NumericQuery)
 
     def test_parse_flex_type_query(self):
         q = self.qfs(["some_float_field:2..3"])
-        self.assertIsInstance(q.subqueries[0], dbcore.query.NumericQuery)
+        self.assertIsInstance(q.subqueries[0], dbcore.query.OrQuery)
 
     def test_empty_query_part(self):
         q = self.qfs([""])

test/test_query.py

@@ -1152,6 +1152,26 @@ def test_get_albums_filter_by_album_flex(self):
         results = self.lib.albums(q)
         self.assert_albums_matched(results, ["Album1"])
 
+    def test_get_albums_filter_by_track_flex(self):
+        q = "item_flex1:Album1"
+        results = self.lib.albums(q)
+        self.assert_albums_matched(results, ["Album1"])
+
+    def test_get_items_filter_by_album_flex(self):
+        q = "album_flex:Album1"
+        results = self.lib.items(q)
+        self.assert_items_matched(results, ["Album1 Item1", "Album1 Item2"])
+
+    def test_filter_by_flex(self):
+        q = "item_flex1:'Item1 Flex1'"
+        results = self.lib.items(q)
+        self.assert_items_matched(results, ["Album1 Item1", "Album2 Item1"])
+
+    def test_filter_by_many_flex(self):
+        q = "item_flex1:'Item1 Flex1' item_flex2:Album1"
+        results = self.lib.items(q)
+        self.assert_items_matched(results, ["Album1 Item1"])
+
 
 def suite():
     return unittest.TestLoader().loadTestsFromName(__name__)