ckanext-search-tweaks 0.6.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

ckanext/search_tweaks/query_relevance/score.py

@@ -1,20 +1,5 @@
-from typing import Optional, Type
-import ckan.plugins.toolkit as tk
-from .storage import (
-    PermanentRedisScoreStorage,
-    DailyRedisScoreStorage,
-    ScoreStorage,
-)
-
-_backends = {
-    "redis-permanent": PermanentRedisScoreStorage,
-    "redis-daily": DailyRedisScoreStorage,
-}
-
-CONFIG_BACKEND = "ckanext.search_tweaks.query_relevance.backend"
-DEFAULT_BACKEND = "redis-daily"
-
-DEFAULT_SCORE_STORAGE_CLASS = DailyRedisScoreStorage
+from .storage import QueryHitTracker
+from .config import get_max_boost_count
 
 
 def normalize_query(query: str) -> str:
@@ -26,50 +11,32 @@ def normalize_query(query: str) -> str:
 
 
 class QueryScore:
-    storage_class: Type[ScoreStorage]
-
-    def __init__(
-        self,
-        id_: str,
-        query: str,
-        *,
-        normalize: bool = True,
-        storage_class: Optional[Type[ScoreStorage]] = None,
-    ):
+    def __init__(self, entity_id: str, query: str, normalize: bool = True):
         if normalize:
             query = normalize_query(query)
 
-        if storage_class:
-            self.storage_class = storage_class
-        else:
-            self.storage_class = self.default_storage_class()
-        self.storage = self.storage_class(id_, query)
+        self.entity_id = entity_id
+        self.query = query
+
+        self.storage = QueryHitTracker(self.entity_id, self.query)
 
     def __int__(self):
         return self.storage.get()
 
-    @staticmethod
-    def default_storage_class() -> Type[ScoreStorage]:
-        return _backends[tk.config.get(CONFIG_BACKEND, DEFAULT_BACKEND)]
-
-    @property
-    def query(self):
-        return self.storage.query
-
-    def increase(self, n: int) -> None:
-        self.storage.inc(n)
-
-    def align(self):
-        self.storage.align()
+    def increase(self, amount: int) -> None:
+        self.storage.increase(amount)
 
     def reset(self):
-        self.storage.reset()
+        self.storage.reset(self.query)
+
+    @classmethod
+    def get_for_query(cls, query: str, limit: int | None = None) -> list[tuple[bytes, float]]:
+        return QueryHitTracker.top(query, limit or get_max_boost_count())
 
     @classmethod
     def get_all(cls):
-        storage = cls.default_storage_class()
-        return storage.scan()
+        return QueryHitTracker.get_all()
 
     @classmethod
-    def get_for(cls, id_: str):
-        return cls.default_storage_class().scan(id_)
+    def reset_all(cls):
+        return QueryHitTracker.reset_all()

ckanext/search_tweaks/query_relevance/storage.py

@@ -1,174 +1,116 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
-from datetime import date, timedelta
-from typing import Any, Iterable, Tuple, cast
+from ckan.lib.redis import connect_to_redis
 
-import ckan.plugins.toolkit as tk
-from ckan.lib.redis import Redis, connect_to_redis
 
-CONFIG_DAILY_AGE = "ckanext.search_tweaks.query_relevance.daily.age"
-DEFAULT_DAILY_AGE = 90
+ScanItem = tuple[str, int]
 
-ScanItem = Tuple[str, str, int]
 
-
-class ScoreStorage(ABC):
-    id: str
-    query: str
-
-    def __init__(self, id_: str, query: str):
-        self.id = id_
+class QueryHitTracker:
+    def __init__(self, entity_id: str, query: str, ttl: int | None = None):
+        self.entity_id = entity_id
         self.query = query
-
-    @abstractmethod
-    def get(self) -> int:
-        """Get current value."""
-        ...
-
-    @abstractmethod
-    def inc(self, by: int) -> None:
-        """Increase current value by the given value."""
-        ...
-
-    @abstractmethod
-    def set(self, value: int) -> None:
-        """Replace current value with the given one."""
-        ...
+        self.conn = self.make_connection()
+        self.ttl = ttl
 
     @classmethod
-    @abstractmethod
-    def scan(cls, id_: str | None = None) -> Iterable[ScanItem]:
-        """Get all the scores."""
-        ...
+    def make_connection(cls):
+        return connect_to_redis()
 
     @classmethod
-    @abstractmethod
-    def reset_storage(cls):
-        """Remove everything from storage."""
-        ...
+    def _key(cls, query: str) -> str:
+        return f"search-tweaks:query-relevance:{query}"
 
-    def reset(self) -> None:
-        """Set current value to zero."""
-        self.set(0)
+    def increase(self, amount: int) -> None:
+        """Increase the score for the specific entity_id + query.
 
-    def align(self) -> None:
-        """Make some cleanup in order to maintain fast and correct value."""
-        return
+        Args:
+            amount: amount to increase the score by
+        """
+        key = self._key(self.query)
+        pipe = self.conn.pipeline()
 
+        pipe.zincrby(key, amount, self.entity_id)
 
-class RedisScoreStorage(ScoreStorage):
-    _conn: Redis | None = None
+        if self.ttl is not None:
+            pipe.expire(key, self.ttl)
 
-    @property
-    def conn(self):
-        if not self._conn:
-            self._conn = self.connect()
-        return self._conn
+        pipe.execute()
 
-    @staticmethod
-    def connect():
-        return connect_to_redis()
+    def get(self) -> int | None:
+        """Get the score for the specific entity_id + query.
 
-    @staticmethod
-    def _common_key_part() -> str:
-        site_id = tk.config["ckan.site_id"]  # type: ignore
-        return f"{site_id}:land:query_scores"
+        Returns:
+            score or None if not found
+        """
+        result = self.conn.zscore(self._key(self.query), self.entity_id)
+        return int(result) if result else 0  # type: ignore
 
     @classmethod
-    def reset_storage(cls):
-        conn = cls.connect()
-        for key in conn.keys(f"{cls._common_key_part()}:*"):
-            conn.delete(key)
-
-    @abstractmethod
-    def _key(self) -> str: ...
-
-    def reset(self):
-        self.conn.delete(self._key())
-
+    def get_all(cls) -> list[tuple[str, str, int]]:
+        """Get all scores.
 
-class PermanentRedisScoreStorage(RedisScoreStorage):
-    """Put all the points into the same cell.
+        Returns:
+            list of (entity_id, query, score) tuples
+        """
+        conn = cls.make_connection()
+        cursor = 0
+        results: list[tuple[str, str, int]] = []
 
-    Sparingly uses memory and must be prefered when there are no extra
-    requirements for invalidation of stats.
+        while True:
+            cursor, keys = conn.scan(cursor=cursor, match=cls._key("*"), count=1000)  # type: ignore
 
-    """
+            for key in keys:
+                query = key.decode().rsplit(":", 1)[-1]
 
-    def set(self, value: int) -> None:
-        self.conn.hset(self._key(), self.query, value)
+                for entity_id, score in conn.zrange(key, 0, -1, withscores=True):  # type: ignore
+                    results.append((entity_id.decode(), query, int(score)))
 
-    def get(self) -> int:
-        return int(self.conn.hget(self._key(), self.query) or 0)
+            if cursor == 0:
+                break
 
-    def inc(self, by: int) -> None:
-        self.conn.hincrby(self._key(), self.query, by)
-
-    def _key(self):
-        return f"{self._common_key_part()}:{self.id}"
+        return results
 
     @classmethod
-    def scan(cls, id_: str | None = None) -> Iterable[ScanItem]:
-        conn = cls.connect()
-        common_key = cls._common_key_part()
-        pattern = f"{common_key}:{id_}" if id_ else f"{common_key}:*"
-        for key in conn.keys(pattern):
-            _, row_id = key.rsplit(b":", 1)
-            for query, score in conn.hgetall(key).items():
-                yield row_id.decode(), query.decode(), int(score)
-
-
-class DailyRedisScoreStorage(RedisScoreStorage):
-    """Store data inside different cells depending on current date.
-
-    The longer index exists, the more memory it consumes. But it can be aligned
-    periodically in order to free memory.
-
-    """
+    def top(cls, query: str, limit: int = 100) -> list[tuple[bytes, float]]:
+        """Return the top N entities for the given query.
 
-    def set(self, value: int) -> None:
-        key = self._key()
-        zkey = self._zkey()
+        Args:
+            query: search query
+            limit (optional): maximum number of entities to return
 
-        self.conn.zadd(key, {zkey: value})
+        Returns:
+            list of (entity_id, score) tuples
+        """
+        conn = cls.make_connection()
 
-    def get(self) -> int:
-        key = self._key()
-        values = self.conn.zrange(key, 0, -1, withscores=True)
-        return self._total(values)
+        return conn.zrevrange(  # type: ignore
+            cls._key(query),
+            0,
+            limit - 1,
+            withscores=True,
+        )
 
-    @staticmethod
-    def _total(values: list[tuple[Any, Any]]) -> int:
-        return int(sum(map(lambda pair: cast(float, pair[1]), values)))
-
-    def inc(self, by: int) -> None:
-        key = self._key()
-        zkey = self._zkey()
-        # type-stubs don't know that signature is (key, amount, value)
-        self.conn.zincrby(key, by, zkey)  # type: ignore
+    @classmethod
+    def reset(cls, query: str) -> None:
+        """Reset scores for the given query.
 
-    def align(self):
-        age = tk.asint(tk.config.get(CONFIG_DAILY_AGE, DEFAULT_DAILY_AGE))
-        verge = bytes((date.today() - timedelta(days=age)).isoformat(), "utf8")
-        key = self._key()
+        Args:
+            query: search query
+        """
+        cls.make_connection().delete(cls._key(query))
 
-        for day in self.conn.zrange(key, 0, -1):
-            if day >= verge:
-                continue
-            self.conn.zrem(key, day)
+    @classmethod
+    def reset_all(cls) -> None:
+        """Reset all scores."""
+        cursor = 0
+        conn = cls.make_connection()
 
-    def _key(self) -> str:
-        return f"{self._common_key_part()}:{self.id}:{self.query}"
+        while True:
+            cursor, keys = conn.scan(cursor=cursor, match=cls._key("*"), count=1000)  # type: ignore
 
-    def _zkey(self):
-        return date.today().isoformat()
+            if keys:
+                conn.delete(*keys)
 
-    @classmethod
-    def scan(cls, id_: str | None = None) -> Iterable[ScanItem]:
-        conn = cls.connect()
-        common_key = cls._common_key_part()
-        pattern = f"{common_key}:{id_}:*" if id_ else f"{common_key}:*"
-        for key in conn.keys(pattern):
-            _, id_, query = key.decode().rsplit(":", 2)
-            yield id_, query, cls(id_, query).get()
+            if cursor == 0:
+                break

ckanext/search_tweaks/spellcheck/helpers.py

@@ -108,7 +108,7 @@ class Collation:
         self.hits = data["hits"]
         self.query = data["collationQuery"]
         changes = data["misspellingsAndCorrections"]
-        self.corrections = dict(zip(changes[::2], changes[1::2]))
+        self.corrections = dict(zip(changes[::2], changes[1::2], strict=True))
 
     def __eq__(self, other):
         if isinstance(other, int):
@@ -145,7 +145,11 @@ class SpellcheckResult:
     def __init__(self, collations: list[Any], suggestions: list[Any]):
         self.collations = [Collation(item) for item in collations[1::2]]
         self.suggestions = dict(
-            zip(suggestions[::2], [s["suggestion"] for s in suggestions[1::2]]),
+            zip(
+                suggestions[::2],
+                [s["suggestion"] for s in suggestions[1::2]],
+                strict=True,
+            ),
         )
 
     def best_collations(self, n: int | None = None) -> list[Collation]:

ckanext/search_tweaks/tests/conftest.py

@@ -3,10 +3,13 @@ import types
 from unittest import mock
 from typing import cast
 
+import factory
 import pytest
+from pytest_factoryboy import register
 
 import ckan.lib.search.query as query
 from ckan.tests.helpers import call_action
+from ckan.tests import factories
 
 
 @pytest.fixture
@@ -25,3 +28,13 @@ def search(monkeypatch):
         return patch.call_args.args[1]
 
     return expose_args
+
+
+@register(_name="dataset")
+class DatasetFactory(factories.Dataset):
+    owner_org = factory.LazyFunction(lambda: OrganizationFactory()["id"])
+
+
+@register(_name="organization")
+class OrganizationFactory(factories.Organization):
+    pass

ckanext/search_tweaks/tests/query_relevance/test_plugin.py

@@ -5,7 +5,7 @@ import ckan.model as model
 import ckanext.search_tweaks.query_relevance as relevance
 
 
-@pytest.mark.usefixtures("with_request_context")
+@pytest.mark.usefixtures("clean_db", "with_request_context")
 class TestPathHasScore:
     @pytest.mark.parametrize(
         "path, has_score",
@@ -17,7 +17,7 @@ class TestPathHasScore:
     )
     def test_search_referrer(self, path, has_score):
         pkg = model.Package(type="dataset")
-        assert relevance._path_has_score_for(path, pkg) is has_score
+        assert relevance._is_scoring_enabled_for_path(path, pkg) is has_score
 
     @pytest.mark.parametrize(
         "path, has_score",
@@ -34,7 +34,7 @@ class TestPathHasScore:
             "get",
             lambda _: model.Group(name="valid", type="organization"),
         )
-        assert relevance._path_has_score_for(path, pkg) is has_score
+        assert relevance._is_scoring_enabled_for_path(path, pkg) is has_score
 
     @pytest.mark.parametrize(
         "path, has_score",
@@ -47,10 +47,10 @@ class TestPathHasScore:
     def test_group_referrer(self, path, has_score, monkeypatch):
         pkg = model.Package(type="dataset")
         monkeypatch.setattr(model.Group, "get", lambda _: model.Group(name="valid"))
-        assert relevance._path_has_score_for(path, pkg) is has_score
+        assert relevance._is_scoring_enabled_for_path(path, pkg) is has_score
 
 
-@pytest.mark.usefixtures("with_request_context")
+@pytest.mark.usefixtures("clean_db", "with_request_context")
 class TestUpdateScore:
     @pytest.mark.parametrize(
         "url, repeat, value",

ckanext/search_tweaks/tests/query_relevance/test_search.py

@@ -0,0 +1,84 @@
+import pytest
+
+from ckan.tests.helpers import call_action
+
+from ckanext.search_tweaks.query_relevance import QueryScore
+
+
+@pytest.mark.usefixtures("with_plugins", "clean_db", "clean_redis", "clean_index")
+class TestSearchScoreBoost:
+    def test_no_score_boost(self, dataset_factory):
+        dataset_factory(title="ocean water")
+        dataset_factory(title="water basin")
+
+        result = call_action("package_search", q="water", fl="id,title,score")[
+            "results"
+        ]
+
+        assert abs(result[0]["score"] - result[1]["score"]) < 0.01
+
+    def test_query_relevance_disabled(self, dataset_factory):
+        dataset_1 = dataset_factory(title="ocean water")
+        dataset_2 = dataset_factory(title="water basin")
+
+        QueryScore(dataset_1["id"], "water").increase(10)
+        QueryScore(dataset_2["id"], "water").increase(5)
+
+        result = call_action(
+            "package_search",
+            q="water",
+            fl="id,title,score",
+            extras={"ext_search_tweaks_disable_query_boost": True},
+        )["results"]
+
+        assert abs(result[0]["score"] - result[1]["score"]) < 0.01
+
+    @pytest.mark.ckan_config("ckanext.search_tweaks.common.prefer_boost", "false")
+    def test_query_relevance_boosted_with_bf(self, dataset_factory):
+        dataset_1 = dataset_factory(title="ocean water")
+        dataset_2 = dataset_factory(title="water basin")
+
+        QueryScore(dataset_1["id"], "water").increase(10)
+        QueryScore(dataset_2["id"], "water").increase(5)
+
+        result = call_action("package_search", q="water", fl="id,title,score")[
+            "results"
+        ]
+
+        assert abs(result[0]["score"] - result[1]["score"]) > 0.01
+
+    @pytest.mark.ckan_config("ckanext.search_tweaks.common.prefer_boost", "true")
+    def test_query_relevance_boosted_with_boost(self, dataset_factory):
+        dataset_1 = dataset_factory(title="ocean water")
+        dataset_2 = dataset_factory(title="water basin")
+
+        QueryScore(dataset_1["id"], "water").increase(10)
+        QueryScore(dataset_2["id"], "water").increase(5)
+
+        result = call_action("package_search", q="water", fl="id,title,score")[
+            "results"
+        ]
+
+        assert abs(result[0]["score"] - result[1]["score"]) > 0.01
+
+    @pytest.mark.skip(reason="use only for profiling")
+    def test_profile_boost_function(self, dataset_factory):
+        first_query = None
+
+        for _ in range(1000):
+            dataset = dataset_factory()
+            ds_query = dataset["title"].split()[0]
+
+            if first_query is None:
+                first_query = ds_query
+
+            QueryScore(dataset["id"], ds_query).increase(1)
+
+        import timeit
+
+        def time_query():
+            call_action("package_search", q=first_query, fl="id,title,score")
+
+        time = timeit.timeit(time_query, number=20)
+
+        print(f"Time: {time}")

ckanext/search_tweaks/tests/query_relevance/test_storage.py

@@ -1,114 +1,38 @@
 import pytest
-from ckanext.search_tweaks.query_relevance.storage import (
-    DailyRedisScoreStorage,
-    PermanentRedisScoreStorage,
-)
 
+from ckanext.search_tweaks.query_relevance.storage import QueryHitTracker
 
-@pytest.fixture
-def storage(storage_class):
-    storage_class.reset_storage()
-    return storage_class
 
+@pytest.mark.usefixtures("clean_redis")
+class TestQueryHitTracker:
+    def test_increase(self):
+        QueryHitTracker("id-1", "hello").increase(1)
+        QueryHitTracker("id-2", "hello").increase(5)
 
-@pytest.mark.parametrize(
-    "storage_class",
-    [
-        PermanentRedisScoreStorage,
-        DailyRedisScoreStorage,
-    ],
-)
-class TestStorages:
-    def test_scan(self, storage):
-        assert list(storage.scan()) == []
+        result = QueryHitTracker.top("hello", 2)
 
-        s1 = storage("key", "query")
-        s1.inc(10)
-        assert sorted(list(storage.scan())) == sorted(
-            [
-                ("key", "query", 10),
-            ],
-        )
+        assert result == [(b"id-2", 5.0), (b"id-1", 1.0)]
 
-        s2 = storage("second key", "second query")
-        s2.inc(5)
-        s1.inc(90)
-        assert sorted(list(storage.scan())) == sorted(
-            [
-                ("key", "query", 100),
-                ("second key", "second query", 5),
-            ],
-        )
+    def test_missing_query(self):
+        result = QueryHitTracker.top("hello")
 
-        s3 = storage("key", "extra query")
-        s3.inc(1)
-        assert sorted(list(storage.scan())) == sorted(
-            [
-                ("key", "query", 100),
-                ("key", "extra query", 1),
-                ("second key", "second query", 5),
-            ],
-        )
+        assert result == []
 
-        assert sorted(list(storage.scan("key"))) == sorted(
-            [
-                ("key", "query", 100),
-                ("key", "extra query", 1),
-            ],
-        )
+    def test_expiration(self):
+        tracker = QueryHitTracker("id-1", "hello", ttl=1)
+        tracker.increase(1)
 
-    def test_missing_key(self, storage):
-        s = storage("not a real key", "not a real query")
-        assert s.get() == 0
+        import time
 
-    def test_set_and_reset(self, storage):
-        s = storage("real key", "real value")
-        s.set(10)
-        assert s.get() == 10
-        s.reset()
-        assert s.get() == 0
+        time.sleep(2)
 
-    def test_increases(self, storage):
-        s1 = storage("real key", "hello")
-        s2 = storage("real key", "world")
+        assert tracker.top("hello") == []
 
-        s1.inc(1)
-        s2.inc(1)
-        s1.inc(1)
-        assert s1.get() == 2
-        assert s2.get() == 1
+    def test_reset(self):
+        tracker = QueryHitTracker("id-1", "hello")
 
+        tracker.increase(1)
+        assert tracker.top("hello") == [(b"id-1", 1.0)]
 
-class TestDailyStorage:
-    @pytest.fixture(autouse=True)
-    def reset_storage(self):
-        DailyRedisScoreStorage.reset_storage()
-
-    def test_score_aggregated(self, freezer):
-        s = DailyRedisScoreStorage("key", "query")
-        freezer.move_to("2012-01-01")
-        s.inc(2)
-        assert s.get() == 2
-
-        freezer.move_to("2012-02-10")
-        s.inc(1)
-        assert s.get() == 3
-
-        freezer.move_to("2012-03-26")
-        s.inc(2)
-        assert s.get() == 5
-
-    def test_score_aligned(self, freezer):
-        s = DailyRedisScoreStorage("key", "query")
-        freezer.move_to("2010-01-01")
-        s.inc(2)
-        freezer.move_to("2011-01-01")
-        s.inc(2)
-
-        freezer.move_to("2012-02-10")
-        s.inc(1)
-        freezer.move_to("2012-03-26")
-        s.inc(2)
-        assert s.get() == 7
-        s.align()
-        assert s.get() == 3
+        tracker.reset("hello")
+        assert tracker.top("hello") == []