nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/external_index_providers/settings.py

@@ -0,0 +1,52 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+
+class ExternalIndexProvidersSettings(BaseSettings):
+    pinecone_upsert_parallelism: int = Field(
+        default=3,
+        title="Pinecone upsert parallelism",
+        description="Number of parallel upserts to Pinecone on each set resource operation",
+    )
+    pinecone_delete_parallelism: int = Field(
+        default=2,
+        title="Pinecone delete parallelism",
+        description="Number of parallel deletes to Pinecone on each delete resource operation",
+    )
+    pinecone_upsert_timeout: float = Field(
+        default=10.0,
+        title="Pinecone upsert timeout",
+        description="Timeout in seconds for each upsert operation to Pinecone",
+    )
+    pinecone_delete_timeout: float = Field(
+        default=10.0,
+        title="Pinecone delete timeout",
+        description="Timeout in seconds for each delete operation to Pinecone",
+    )
+    pinecone_query_timeout: float = Field(
+        default=10.0,
+        title="Pinecone query timeout",
+        description="Timeout in seconds for each query operation to Pinecone",
+    )
+
+
+settings = ExternalIndexProvidersSettings()

nucliadb/common/http_clients/auth.py

@@ -44,14 +44,11 @@ class NucliaAuthHTTPClient:
     def __init__(self):
         self.session = aiohttp.ClientSession()
         self.base_url = (
-            nuclia_settings.nuclia_public_url.format(zone=nuclia_settings.nuclia_zone)
-            + "/api"
+            nuclia_settings.nuclia_public_url.format(zone=nuclia_settings.nuclia_zone) + "/api"
         )
         self.headers = {}
         if nuclia_settings.nuclia_service_account is not None:
-            self.headers[
-                "X-NUCLIA-NUAKEY"
-            ] = f"Bearer {nuclia_settings.nuclia_service_account}"
+            self.headers["X-NUCLIA-NUAKEY"] = f"Bearer {nuclia_settings.nuclia_service_account}"
 
     async def __aenter__(self):
         return self
@@ -67,4 +64,4 @@ class NucliaAuthHTTPClient:
         async with self.session.get(url, headers=self.headers) as resp:
             resp_text = await resp.text()
             check_status(resp, resp_text)
-            return AuthInfoResponse.parse_raw(resp_text)
+            return AuthInfoResponse.model_validate_json(resp_text)

nucliadb/common/http_clients/processing.py

@@ -48,10 +48,7 @@ def get_processing_api_url() -> str:
             + "/api/v1/processing"
         )
     else:
-        return (
-            nuclia_settings.nuclia_processing_cluster_url
-            + "/api/v1/internal/processing"
-        )
+        return nuclia_settings.nuclia_processing_cluster_url + "/api/v1/internal/processing"
 
 
 class PullResponse(pydantic.BaseModel):
@@ -159,9 +156,7 @@ class ProcessingHTTPClient:
         self.base_url = get_processing_api_url()
         self.headers = {}
         if nuclia_settings.nuclia_service_account is not None:
-            self.headers[
-                "X-STF-NUAKEY"
-            ] = f"Bearer {nuclia_settings.nuclia_service_account}"
+            self.headers["X-STF-NUAKEY"] = f"Bearer {nuclia_settings.nuclia_service_account}"
 
     async def __aenter__(self):
         return self
@@ -187,7 +182,7 @@ class ProcessingHTTPClient:
         async with self.session.get(url, headers=self.headers, params=params) as resp:
             resp_text = await resp.text()
             check_status(resp, resp_text)
-            return PullResponse.parse_raw(resp_text)
+            return PullResponse.model_validate_json(resp_text)
 
     async def pull_position(self, partition: str) -> int:
         url = self.base_url + "/pull/position"
@@ -195,7 +190,7 @@ class ProcessingHTTPClient:
         async with self.session.get(url, headers=self.headers, params=params) as resp:
             resp_text = await resp.text()
             check_status(resp, resp_text)
-            data = PullPosition.parse_raw(resp_text)
+            data = PullPosition.model_validate_json(resp_text)
             return data.cursor
 
     async def requests(
@@ -217,7 +212,7 @@ class ProcessingHTTPClient:
         async with self.session.get(url, headers=self.headers, params=params) as resp:
             resp_text = await resp.text()
             check_status(resp, resp_text)
-            return RequestsResults.parse_raw(resp_text)
+            return RequestsResults.model_validate_json(resp_text)
 
     async def stats(self, kbid: str, timeout: Optional[float] = 1.0) -> StatsResponse:
         url = self.base_url + "/stats"
@@ -229,4 +224,4 @@ class ProcessingHTTPClient:
         ) as resp:
             resp_text = await resp.text()
             check_status(resp, resp_text)
-            return StatsResponse.parse_raw(resp_text)
+            return StatsResponse.model_validate_json(resp_text)

nucliadb/common/http_clients/utils.py

@@ -30,9 +30,7 @@ def check_status(resp: aiohttp.ClientResponse, resp_text: str) -> None:
     elif resp.status == 404:
         raise exceptions.NotFoundException(f"Resource not found: {resp_text}")
     elif resp.status in (401, 403):
-        raise exceptions.AuthorizationException(
-            f"Unauthorized to access: {resp.status}"
-        )
+        raise exceptions.AuthorizationException(f"Unauthorized to access: {resp.status}")
     elif resp.status == 429:
         raise exceptions.RateLimitException("Rate limited")
     else:

nucliadb/common/ids.py

@@ -0,0 +1,240 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+
+"""
+This module aims to centralize how we build ids for resources, fields,
+paragraphs... Avoiding spread of id construction and parsing everywhere
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from nucliadb_protos.resources_pb2 import FieldType
+
+FIELD_TYPE_STR_TO_PB: dict[str, FieldType.ValueType] = {
+    "t": FieldType.TEXT,
+    "f": FieldType.FILE,
+    "u": FieldType.LINK,
+    "a": FieldType.GENERIC,
+    "c": FieldType.CONVERSATION,
+}
+
+FIELD_TYPE_PB_TO_STR = {v: k for k, v in FIELD_TYPE_STR_TO_PB.items()}
+
+
+@dataclass
+class FieldId:
+    """
+    Field ids are used to identify fields in resources. They usually have the following format:
+
+        `rid/field_type/field_key`
+
+    where field type is one of: `t`, `f`, `u`, `a`, `c` (text, file, link, generic, conversation)
+    and field_key is an identifier for that field type on the resource, usually chosen by the user.
+
+    In some cases, fields can have subfields, for example, in conversations, where each part of the
+    conversation is a subfield. In those cases, the id has the following format:
+
+        `rid/field_type/field_key/subfield_id`
+
+    Examples:
+
+    >>> FieldId(rid="rid", type="u", key="/my-link")
+    FieldID("rid/u/my-link")
+    >>> FieldId.from_string("rid/u/my-link")
+    FieldID("rid/u/my-link")
+    """
+
+    rid: str
+    type: str
+    key: str
+    # also knwon as `split`, this indicates a part of a field in, for example, conversations
+    subfield_id: Optional[str] = None
+
+    def __repr__(self) -> str:
+        return f"FieldId({self.full()})"
+
+    def short_without_subfield(self) -> str:
+        return f"/{self.type}/{self.key}"
+
+    def full(self) -> str:
+        if self.subfield_id is None:
+            return f"{self.rid}/{self.type}/{self.key}"
+        else:
+            return f"{self.rid}/{self.type}/{self.key}/{self.subfield_id}"
+
+    def __hash__(self) -> int:
+        return hash(self.full())
+
+    @property
+    def pb_type(self) -> FieldType.ValueType:
+        return FIELD_TYPE_STR_TO_PB[self.type]
+
+    @classmethod
+    def from_pb(
+        cls, rid: str, field_type: FieldType.ValueType, key: str, subfield_id: Optional[str] = None
+    ) -> "FieldId":
+        return cls(rid=rid, type=FIELD_TYPE_PB_TO_STR[field_type], key=key, subfield_id=subfield_id)
+
+    @classmethod
+    def from_string(cls, value: str) -> "FieldId":
+        """
+        Parse a FieldId from a string
+        Example:
+        >>> fid = FieldId.from_string("rid/u/foo")
+        >>> fid
+        FieldId("rid/u/foo")
+        >>> fid.type
+        'u'
+        >>> fid.key
+        'foo'
+        >>> FieldId.from_string("rid/u/foo/subfield_id").subfield_id
+        'subfield_id'
+        """
+        parts = value.split("/")
+        if len(parts) == 3:
+            rid, _type, key = parts
+            if _type not in FIELD_TYPE_STR_TO_PB:
+                raise ValueError(f"Invalid FieldId: {value}")
+            return cls(rid=rid, type=_type, key=key)
+        elif len(parts) == 4:
+            rid, _type, key, subfield_id = parts
+            if _type not in FIELD_TYPE_STR_TO_PB:
+                raise ValueError(f"Invalid FieldId: {value}")
+            return cls(
+                rid=rid,
+                type=_type,
+                key=key,
+                subfield_id=subfield_id,
+            )
+        else:
+            raise ValueError(f"Invalid FieldId: {value}")
+
+
+@dataclass
+class ParagraphId:
+    field_id: FieldId
+    paragraph_start: int
+    paragraph_end: int
+
+    def __repr__(self) -> str:
+        return f"ParagraphId({self.full()})"
+
+    def full(self) -> str:
+        return f"{self.field_id.full()}/{self.paragraph_start}-{self.paragraph_end}"
+
+    def __hash__(self) -> int:
+        return hash(self.full())
+
+    @property
+    def rid(self) -> str:
+        return self.field_id.rid
+
+    @classmethod
+    def from_string(cls, value: str) -> "ParagraphId":
+        parts = value.split("/")
+        paragraph_range = parts[-1]
+        start, end = map(int, paragraph_range.split("-"))
+        field_id = FieldId.from_string("/".join(parts[:-1]))
+        return cls(field_id=field_id, paragraph_start=start, paragraph_end=end)
+
+    @classmethod
+    def from_vector_id(cls, vid: "VectorId") -> "ParagraphId":
+        """
+        Returns a ParagraphId from a vector_key (the index part of the vector_key is ignored).
+        >>> vid = VectorId.from_string("rid/u/field_id/0/0-1")
+        >>> ParagraphId.from_vector_id(vid)
+        ParagraphId("rid/u/field_id/0-1")
+        """
+        return cls(
+            field_id=vid.field_id,
+            paragraph_start=vid.vector_start,
+            paragraph_end=vid.vector_end,
+        )
+
+
+@dataclass
+class VectorId:
+    """
+    Ids of vectors are very similar to ParagraphIds, but for legacy reasons, they have an index
+    indicating the position of the corresponding text block in the list of text blocks for the field.
+
+    Examples:
+
+    >>> VectorId.from_string("rid/u/field_id/0/0-10")
+    VectorId("rid/u/field_id/0/0-10")
+    >>> VectorId(
+    ...    field_id=FieldId.from_string("rid/u/field_id"),
+    ...    index=0,
+    ...    vector_start=0,
+    ...    vector_end=10,
+    ... )
+    VectorId("rid/u/field_id/0/0-10")
+    """
+
+    field_id: FieldId
+    index: int
+    vector_start: int
+    vector_end: int
+
+    def __repr__(self) -> str:
+        return f"VectorId({self.full()})"
+
+    def full(self) -> str:
+        return f"{self.field_id.full()}/{self.index}/{self.vector_start}-{self.vector_end}"
+
+    def __hash__(self) -> int:
+        return hash(self.full())
+
+    @property
+    def rid(self) -> str:
+        return self.field_id.rid
+
+    @classmethod
+    def from_string(cls, value: str) -> "VectorId":
+        parts = value.split("/")
+        vector_range = parts[-1]
+        start, end = map(int, vector_range.split("-"))
+        index = int(parts[-2])
+        field_id = FieldId.from_string("/".join(parts[:-2]))
+        return cls(field_id=field_id, index=index, vector_start=start, vector_end=end)
+
+
+def extract_data_augmentation_id(generated_field_id: str) -> Optional[str]:
+    """Data augmentation generated fields have a strict id with the following
+    format:
+    `da-{task_id}-{original:field_type}-{original:field_id}[-{original:split}]`
+
+    @return the `task_id`
+
+    ATENTION: we are assuming ids have been properly generated and `-` is not a
+    valid character, otherwise, this extraction would be wrong and a partial id
+    would be returned.
+
+    """
+    parts = generated_field_id.split("-")
+
+    if len(parts) < 4:
+        return None
+
+    if parts[0] != "da":
+        return None
+
+    return parts[1] or None

nucliadb/common/locking.py

@@ -35,11 +35,15 @@ logger = logging.getLogger(__name__)
 
 NEW_SHARD_LOCK = "new-shard-{kbid}"
 RESOURCE_INDEX_LOCK = "resource-index-{kbid}-{resource_id}"
+RESOURCE_CREATION_SLUG_LOCK = "resource-creation-{kbid}-{resource_slug}"
 KB_SHARDS_LOCK = "shards-kb-{kbid}"
+MIGRATIONS_LOCK = "migration"
 
 
 class ResourceLocked(Exception):
-    ...
+    def __init__(self, key: str):
+        self.key = key
+        super().__init__(f"{key} is locked")
 
 
 @dataclass
@@ -59,7 +63,8 @@ class _Lock:
         expire_timeout: float,
         refresh_timeout: float,
     ):
-        self.key = "/distributed/locks/" + key
+        self.user_key = key
+        self.key = "/distributed/locks/" + self.user_key
         self.lock_timeout = lock_timeout
         self.expire_timeout = expire_timeout
         self.refresh_timeout = refresh_timeout
@@ -79,40 +84,52 @@ class _Lock:
                     else:
                         if time.time() > lock_data.expires_at:
                             # if current time is greater than when it expires, take it over
-                            await self._set_lock_value(txn)
+                            await self._update_lock_value(txn)
                             await txn.commit()
                             break
 
                         if time.time() > start + self.lock_timeout:
                             # if current time > start time + lock timeout
                             # we've waited too long, raise exception that, we can't get the lock
-                            raise ResourceLocked()
+                            raise ResourceLocked(key=self.user_key)
             except ConflictError:
                 # if we get a conflict error, retry
                 pass
-            await asyncio.sleep(0.1)  # sleep before trying againt
+            await asyncio.sleep(0.1)  # sleep before trying again
         self.task = asyncio.create_task(self._refresh_task())
         return self
 
     async def get_lock_data(self, txn: Transaction) -> Optional[LockValue]:
-        existing_data = await txn.get(self.key)
+        existing_data = await txn.get(self.key, for_update=True)
         if existing_data is None:
             return None
         else:
             return LockValue(**orjson.loads(existing_data))
 
-    async def _set_lock_value(self, txn: Transaction) -> None:
+    async def _update_lock_value(self, txn: Transaction) -> None:
+        """
+        Update the value for the lock.
+        """
         await txn.set(
             self.key,
             orjson.dumps(LockValue(self.value, time.time() + self.expire_timeout)),
         )
 
+    async def _set_lock_value(self, txn: Transaction) -> None:
+        """
+        Set the value for the lock. If lock already exists, it doesn't update and raises a ConflictError.
+        """
+        await txn.insert(
+            self.key,
+            orjson.dumps(LockValue(self.value, time.time() + self.expire_timeout)),
+        )
+
     async def _refresh_task(self) -> None:
         while True:
             try:
                 await asyncio.sleep(self.refresh_timeout)
                 async with self.driver.transaction() as txn:
-                    await self._set_lock_value(txn)
+                    await self._update_lock_value(txn)
                     await txn.commit()
             except (asyncio.CancelledError, RuntimeError):
                 return
@@ -125,21 +142,34 @@ class _Lock:
             await txn.delete(self.key)
             await txn.commit()
 
-    async def is_locked(self, key: str) -> bool:
+    async def is_locked(self) -> bool:
         async with get_driver().transaction(read_only=True) as txn:
             lock_data = await self.get_lock_data(txn)
-        return lock_data is None or time.time() > lock_data.expires_at
+        return lock_data is not None and time.time() < lock_data.expires_at
 
 
 def distributed_lock(
     key: str,
-    lock_timeout: float = 60.0,  # max time to wait for lock
-    expire_timeout: float = 30.0,  # how long by default the lock will be held without a refresh
-    refresh_timeout: float = 10.0,  # how often to refresh
+    lock_timeout: float = 60.0,
+    expire_timeout: float = 30.0,
+    refresh_timeout: float = 10.0,
 ) -> _Lock:
+    """
+    Context manager to get a distributed lock on a key.
+
+    Params:
+    - key: the key to lock with
+    - lock_timeout: maximum time to wait for the lock before ResourceLocked is raised.
+    - expire_timeout: how long by default the lock will be held without a refresh
+    - refresh_timeout: how often to refresh the lock
+    """
     return _Lock(
         key,
         lock_timeout=lock_timeout,
         expire_timeout=expire_timeout,
         refresh_timeout=refresh_timeout,
     )
+
+
+async def is_locked(key: str) -> bool:
+    return await distributed_lock(key).is_locked()

nucliadb/common/maindb/driver.py

@@ -23,7 +23,7 @@ import asyncio
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator, Optional
 
-DEFAULT_SCAN_LIMIT = 10
+DEFAULT_SCAN_LIMIT = -1
 DEFAULT_BATCH_SCAN_LIMIT = 500
 
 
@@ -37,18 +37,24 @@ class Transaction:
     async def commit(self):
         raise NotImplementedError()
 
-    async def batch_get(self, keys: list[str]) -> list[Optional[bytes]]:
+    async def batch_get(self, keys: list[str], for_update: bool = False) -> list[Optional[bytes]]:
         raise NotImplementedError()
 
-    async def get(self, key: str) -> Optional[bytes]:
+    async def get(self, key: str, for_update: bool = False) -> Optional[bytes]:
         raise NotImplementedError()
 
     async def set(self, key: str, value: bytes):
         raise NotImplementedError()
 
+    async def insert(self, key: str, value: bytes):
+        return await self.set(key, value)
+
     async def delete(self, key: str):
         raise NotImplementedError()
 
+    async def delete_by_prefix(self, prefix: str) -> None:
+        raise NotImplementedError()
+
     def keys(
         self, match: str, count: int = DEFAULT_SCAN_LIMIT, include_start: bool = True
     ) -> AsyncGenerator[str, None]:
@@ -74,36 +80,6 @@ class Driver:
                 except Exception:
                     pass
 
-    async def begin(self, read_only: bool = False) -> Transaction:
-        raise NotImplementedError()
-
     @asynccontextmanager
-    async def transaction(
-        self, wait_for_abort: bool = True, read_only: bool = False
-    ) -> AsyncGenerator[Transaction, None]:
-        """
-        Use to make sure transaction is always aborted.
-
-        :param wait_for_abort: If True, wait for abort to finish before returning.
-                               If False, abort is done in background (unless there
-                               is an error)
-        """
-        txn: Optional[Transaction] = None
-        error: bool = False
-        try:
-            txn = await self.begin(read_only=read_only)
-            yield txn
-        except Exception:
-            error = True
-            raise
-        finally:
-            if txn is not None and txn.open:
-                if error or wait_for_abort:
-                    await txn.abort()
-                else:
-                    self._async_abort(txn)
-
-    def _async_abort(self, txn: Transaction):
-        task = asyncio.create_task(txn.abort())
-        task.add_done_callback(lambda task: self._abort_tasks.remove(task))
-        self._abort_tasks.append(task)
+    async def transaction(self, read_only: bool = False) -> AsyncGenerator[Transaction, None]:
+        yield Transaction()

nucliadb/common/maindb/exceptions.py

@@ -17,13 +17,13 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-class ConflictError(Exception):
-    ...
+class ConflictError(Exception): ...
 
 
-class NotFoundError(Exception):
-    ...
+class NotFoundError(Exception): ...
 
 
-class UnsetUtility(Exception):
-    ...
+class UnsetUtility(Exception): ...
+
+
+class MaindbServerError(Exception): ...

nucliadb/common/maindb/local.py

@@ -19,7 +19,8 @@
 #
 import glob
 import os
-from typing import Optional
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, Optional
 
 from nucliadb.common.maindb.driver import (
     DEFAULT_BATCH_SCAN_LIMIT,
@@ -105,7 +106,7 @@ class LocalTransaction(Transaction):
         self.clean()
         self.open = False
 
-    async def batch_get(self, keys: list[str]) -> list[Optional[bytes]]:
+    async def batch_get(self, keys: list[str], for_update: bool = False) -> list[Optional[bytes]]:
         results: list[Optional[bytes]] = []
         for key in keys:
             obj = await self.get(key)
@@ -124,7 +125,7 @@ class LocalTransaction(Transaction):
 
         return results
 
-    async def get(self, key: str) -> Optional[bytes]:
+    async def get(self, key: str, for_update: bool = False) -> Optional[bytes]:
         if key in self.deleted_keys:
             raise KeyError(f"Not found {key}")
 
@@ -159,9 +160,15 @@ class LocalTransaction(Transaction):
         if key in self.modified_keys:
             del self.modified_keys[key]
 
-    async def keys(
-        self, match: str, count: int = DEFAULT_SCAN_LIMIT, include_start: bool = True
-    ):
+    async def delete_by_prefix(self, prefix: str) -> None:
+        keys = []
+        for key in self.modified_keys.keys():
+            if key.startswith(prefix):
+                keys.append(key)
+        for key in keys:
+            await self.delete(key)
+
+    async def keys(self, match: str, count: int = DEFAULT_SCAN_LIMIT, include_start: bool = True):
         prev_key = None
 
         get_all_keys = count == -1
@@ -195,7 +202,7 @@ class LocalTransaction(Transaction):
 
     async def count(self, match: str) -> int:
         value = 0
-        async for _ in self.keys(match, count=-1):
+        async for _ in self.keys(match):
             value += 1
         return value
 
@@ -214,7 +221,13 @@ class LocalDriver(Driver):
     async def finalize(self):
         pass
 
-    async def begin(self, read_only: bool = False) -> LocalTransaction:
+    @asynccontextmanager
+    async def transaction(self, read_only: bool = False) -> AsyncGenerator[Transaction, None]:
         if self.url is None:
             raise AttributeError("Invalid url")
-        return LocalTransaction(self.url, self)
+        txn = LocalTransaction(self.url, self)
+        try:
+            yield txn
+        finally:
+            if txn.open:
+                await txn.abort()