nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/common/cache.py

@@ -24,7 +24,7 @@ from abc import ABC, abstractmethod
 from contextvars import ContextVar
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Generic, Optional, TypeVar
+from typing import Generic, TypeVar
 
 import backoff
 from async_lru import _LRUCacheWrapper, alru_cache
@@ -66,9 +66,9 @@ class Cache(Generic[K, T], ABC):
 
     """
 
-    cache: _LRUCacheWrapper[Optional[T]]
+    cache: _LRUCacheWrapper[T | None]
 
-    async def get(self, *args: K.args, **kwargs: K.kwargs) -> Optional[T]:
+    async def get(self, *args: K.args, **kwargs: K.kwargs) -> T | None:
         result = await self.cache(*args)
         # Do not cache None
         if result is None:
@@ -88,7 +88,7 @@ class Cache(Generic[K, T], ABC):
 class ResourceCache(Cache[[str, str], ResourceORM]):
     def __init__(self, cache_size: int) -> None:
         @alru_cache(maxsize=cache_size)
-        async def _get_resource(kbid: str, rid: str) -> Optional[ResourceORM]:
+        async def _get_resource(kbid: str, rid: str) -> ResourceORM | None:
             storage = await get_storage()
             async with get_driver().ro_transaction() as txn:
                 kb = KnowledgeBoxORM(txn, storage, kbid)
@@ -115,7 +115,7 @@ class ExtractedTextCache(Cache[[str, FieldId], ExtractedText]):
     def __init__(self, cache_size: int) -> None:
         @alru_cache(maxsize=cache_size)
         @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
-        async def _get_extracted_text(kbid: str, field_id: FieldId) -> Optional[ExtractedText]:
+        async def _get_extracted_text(kbid: str, field_id: FieldId) -> ExtractedText | None:
             storage = await get_storage()
             try:
                 sf = storage.file_extracted(
@@ -144,18 +144,18 @@ class ExtractedTextCache(Cache[[str, FieldId], ExtractedText]):
 
 # Global caches (per asyncio task)
 
-rcache: ContextVar[Optional[ResourceCache]] = ContextVar("rcache", default=None)
-etcache: ContextVar[Optional[ExtractedTextCache]] = ContextVar("etcache", default=None)
+rcache: ContextVar[ResourceCache | None] = ContextVar("rcache", default=None)
+etcache: ContextVar[ExtractedTextCache | None] = ContextVar("etcache", default=None)
 
 
 # Cache management
 
 
-def get_resource_cache() -> Optional[ResourceCache]:
+def get_resource_cache() -> ResourceCache | None:
     return rcache.get()
 
 
-def get_extracted_text_cache() -> Optional[ExtractedTextCache]:
+def get_extracted_text_cache() -> ExtractedTextCache | None:
     return etcache.get()
 
 

nucliadb/common/catalog/interface.py

@@ -22,7 +22,7 @@ from __future__ import annotations
 import abc
 import datetime
 from dataclasses import dataclass
-from typing import Literal, Optional, Union
+from typing import Literal
 
 from pydantic import BaseModel, Field
 
@@ -49,22 +49,22 @@ class CatalogResourceData(BaseModel):
 class CatalogExpression:
     @dataclass
     class Date:
-        field: Union[Literal["created_at"], Literal["modified_at"]]
-        since: Optional[datetime.datetime]
-        until: Optional[datetime.datetime]
+        field: Literal["created_at"] | Literal["modified_at"]
+        since: datetime.datetime | None
+        until: datetime.datetime | None
 
-    bool_and: Optional[list["CatalogExpression"]] = None
-    bool_or: Optional[list["CatalogExpression"]] = None
-    bool_not: Optional["CatalogExpression"] = None
-    date: Optional[Date] = None
-    facet: Optional[str] = None
-    resource_id: Optional[str] = None
+    bool_and: list[CatalogExpression] | None = None
+    bool_or: list[CatalogExpression] | None = None
+    bool_not: CatalogExpression | None = None
+    date: Date | None = None
+    facet: str | None = None
+    resource_id: str | None = None
 
 
 class CatalogQuery(BaseModel):
     kbid: str
-    query: Optional[search_models.CatalogQuery] = Field(description="Full-text search query")
-    filters: Optional[CatalogExpression] = Field(description="Filters to apply to the search")
+    query: search_models.CatalogQuery | None = Field(description="Full-text search query")
+    filters: CatalogExpression | None = Field(description="Filters to apply to the search")
     sort: search_models.SortOptions = Field(description="Sorting option")
     faceted: list[str] = Field(description="List of facets to compute during the search")
     page_size: int = Field(description="Used for pagination. Maximum page size is 100")

nucliadb/common/catalog/pg.py

@@ -21,10 +21,11 @@
 import logging
 import re
 from collections import defaultdict
-from typing import Any, Literal, Union, cast
+from typing import Any, Literal, cast
 
 from psycopg import AsyncCursor, sql
 from psycopg.rows import DictRow, dict_row
+from typing_extensions import assert_never
 
 from nucliadb.common.catalog.interface import (
     Catalog,
@@ -267,32 +268,46 @@ async def _faceted_search_unfiltered(
 ):
     facet_params: dict[str, Any] = {}
     facet_sql: sql.Composable
-    if len(tmp_facets) <= 5:
-        # Asking for few facets, strictly filter to what we need in the query
-        prefixes_sql = []
-        for cnt, prefix in enumerate(tmp_facets.keys()):
-            prefixes_sql.append(
-                sql.SQL("(facet LIKE {} AND POSITION('/' IN RIGHT(facet, {})) = 0)").format(
-                    sql.Placeholder(f"facet_{cnt}"), sql.Placeholder(f"facet_len_{cnt}")
+    if list(tmp_facets.keys()) == ["/n/s"]:
+        # Special case when querying only for status. We know the list of possible facets and optimize
+        # by asking for each facet separately which makes better use of the index
+        sqls = []
+        for status in ["PENDING", "PROCESSED", "ERROR", "EMPTY"]:
+            sqls.append(
+                sql.SQL(
+                    "SELECT facet, COUNT(*) FROM catalog_facets WHERE kbid = %(kbid)s AND facet = '/n/s/{}' GROUP BY facet".format(
+                        status
+                    )
                 )
             )
-            facet_params[f"facet_{cnt}"] = f"{prefix}/%"
-            facet_params[f"facet_len_{cnt}"] = -(len(prefix) + 1)
-        facet_sql = sql.SQL("AND {}").format(sql.SQL(" OR ").join(prefixes_sql))
-    elif all((facet.startswith("/l") or facet.startswith("/n/i") for facet in tmp_facets.keys())):
-        # Special case for the catalog query, which can have many facets asked for
-        # Filter for the categories (icon and labels) in the query, filter the rest in the code below
-        facet_sql = sql.SQL("AND (facet LIKE '/l/%%' OR facet like '/n/i/%%')")
+        await cur.execute(sql.SQL(" UNION ").join(sqls), {"kbid": catalog_query.kbid})
     else:
-        # Worst case: ask for all facets and filter here. This is faster than applying lots of filters
-        facet_sql = sql.SQL("")
-
-    await cur.execute(
-        sql.SQL(
-            "SELECT facet, COUNT(*) FROM catalog_facets WHERE kbid = %(kbid)s {} GROUP BY facet"
-        ).format(facet_sql),
-        {"kbid": catalog_query.kbid, **facet_params},
-    )
+        if len(tmp_facets) <= 5:
+            # Asking for few facets, strictly filter to what we need in the query
+            prefixes_sql = []
+            for cnt, prefix in enumerate(tmp_facets.keys()):
+                prefixes_sql.append(
+                    sql.SQL("(facet LIKE {} AND POSITION('/' IN RIGHT(facet, {})) = 0)").format(
+                        sql.Placeholder(f"facet_{cnt}"), sql.Placeholder(f"facet_len_{cnt}")
+                    )
+                )
+                facet_params[f"facet_{cnt}"] = f"{prefix}/%"
+                facet_params[f"facet_len_{cnt}"] = -(len(prefix) + 1)
+            facet_sql = sql.SQL("AND {}").format(sql.SQL(" OR ").join(prefixes_sql))
+        elif all(facet.startswith("/l") or facet.startswith("/n/i") for facet in tmp_facets.keys()):
+            # Special case for the catalog query, which can have many facets asked for
+            # Filter for the categories (icon and labels) in the query, filter the rest in the code below
+            facet_sql = sql.SQL("AND (facet LIKE '/l/%%' OR facet like '/n/i/%%')")
+        else:
+            # Worst case: ask for all facets and filter here. This is faster than applying lots of filters
+            facet_sql = sql.SQL("")
+
+        await cur.execute(
+            sql.SQL(
+                "SELECT facet, COUNT(*) FROM catalog_facets WHERE kbid = %(kbid)s {} GROUP BY facet"
+            ).format(facet_sql),
+            {"kbid": catalog_query.kbid, **facet_params},
+        )
 
     # Only keep the facets we asked for
     for row in await cur.fetchall():
@@ -374,10 +389,7 @@ def _prepare_query_search(query: search_models.CatalogQuery, params: dict[str, A
         params["query"] = "%" + query.query + "%"
         return sql.SQL("title ILIKE %(query)s")
     else:  # pragma: no cover
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
-        return sql.SQL("")
+        assert_never(query.match)
 
 
 def _convert_filter(expr: CatalogExpression, filter_params: dict[str, Any]) -> sql.Composable:
@@ -409,7 +421,7 @@ def _convert_filter(expr: CatalogExpression, filter_params: dict[str, Any]) -> s
 
 def _convert_boolean_op(
     operands: list[CatalogExpression],
-    op: Union[Literal["and"], Literal["or"]],
+    op: Literal["and"] | Literal["or"],
     filter_params: dict[str, Any],
 ) -> sql.Composable:
     array_op = sql.SQL("@>" if op == "and" else "&&")

nucliadb/common/catalog/utils.py

@@ -40,17 +40,17 @@ def build_catalog_resource_data(resource: Resource, index_message: IndexMessage)
     }
 
     # Labels from the resource and classification labels from each field
-    labels = [label for label in index_message.labels]
+    labels = {label for label in index_message.labels}
     for classification in resource.basic.computedmetadata.field_classifications:
         for clf in classification.classifications:
             label = f"/l/{clf.labelset}/{clf.label}"
             if label not in cancelled_labels:
-                labels.append(label)
+                labels.add(label)
 
     return CatalogResourceData(
         title=resource.basic.title,
         created_at=created_at,
         modified_at=modified_at,
-        labels=labels,
+        labels=list(labels),
         slug=resource.basic.slug,
     )

nucliadb/common/cluster/manager.py

@@ -20,7 +20,8 @@
 import asyncio
 import logging
 import uuid
-from typing import Any, Awaitable, Callable, Optional
+from collections.abc import Awaitable, Callable
+from typing import Any
 
 from nidx_protos import noderesources_pb2, nodewriter_pb2
 from nidx_protos.nodewriter_pb2 import (
@@ -96,7 +97,7 @@ class KBShardManager:
     # TODO: move to data manager
     async def get_current_active_shard(
         self, txn: Transaction, kbid: str
-    ) -> Optional[writer_pb2.ShardObject]:
+    ) -> writer_pb2.ShardObject | None:
         kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=False)
         if kb_shards is None:
             return None
@@ -196,7 +197,7 @@ class KBShardManager:
         txid: int,
         partition: str,
         kb: str,
-        reindex_id: Optional[str] = None,
+        reindex_id: str | None = None,
         source: IndexMessageSource.ValueType = IndexMessageSource.PROCESSOR,
     ) -> None:
         """
@@ -306,7 +307,7 @@ class StandaloneKBShardManager(KBShardManager):
         txid: int,
         partition: str,
         kb: str,
-        reindex_id: Optional[str] = None,
+        reindex_id: str | None = None,
         source: IndexMessageSource.ValueType = IndexMessageSource.PROCESSOR,
     ) -> None:
         """