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

nucliadb/common/cluster/discovery/k8s.py

@@ -69,9 +69,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
     ) -> IndexNodeMetadata:
         async with self.update_lock:
             if pod_name not in self.node_id_cache:
-                self.node_id_cache[pod_name] = await self._query_node_metadata(
-                    node_ip, read_replica
-                )
+                self.node_id_cache[pod_name] = await self._query_node_metadata(node_ip, read_replica)
             else:
                 self.node_id_cache[pod_name].address = node_ip
             self.node_id_cache[pod_name].updated_at = time.time()
@@ -84,12 +82,10 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
         This method will update global node state by utilizing the cluster manager
         to add or remove nodes.
         """
-        status: kubernetes_asyncio.client.models.v1_pod_status.V1PodStatus = event[
+        status: kubernetes_asyncio.client.models.v1_pod_status.V1PodStatus = event["object"].status
+        event_metadata: kubernetes_asyncio.client.models.v1_object_meta.V1ObjectMeta = event[
             "object"
-        ].status
-        event_metadata: kubernetes_asyncio.client.models.v1_object_meta.V1ObjectMeta = (
-            event["object"].metadata
-        )
+        ].metadata
 
         ready = status.container_statuses is not None
         if event["type"] == "DELETED":
@@ -137,7 +133,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
         if ready:
             node = manager.get_index_node(node_data.node_id)
             if node is None:
-                logger.warning(
+                logger.info(
                     "Adding node",
                     extra={
                         "node_id": node_data.node_id,
@@ -162,7 +158,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
         else:
             node = manager.get_index_node(node_data.node_id)
             if node is not None:
-                logger.warning(
+                logger.info(
                     f"Remove node",
                     extra={
                         "node_id": node_data.node_id,
@@ -199,9 +195,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
                             except NodeConnectionError:  # pragma: no cover
                                 pass
                             except Exception:  # pragma: no cover
-                                logger.exception(
-                                    "Error while updating node", exc_info=True
-                                )
+                                logger.exception("Error while updating node", exc_info=True)
                     except (
                         asyncio.CancelledError,
                         KeyboardInterrupt,
@@ -259,9 +253,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
                             continue
                         existing = self.node_id_cache[pod_name]
                         try:
-                            self.node_id_cache[
-                                pod_name
-                            ] = await self._query_node_metadata(
+                            self.node_id_cache[pod_name] = await self._query_node_metadata(
                                 existing.address,
                                 read_replica=existing.primary_id is not None,
                             )
@@ -301,9 +293,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
 
     async def initialize(self) -> None:
         self.cluster_task = asyncio.create_task(self.watch_k8s_for_updates())
-        self.update_node_data_cache_task = asyncio.create_task(
-            self.update_node_data_cache()
-        )
+        self.update_node_data_cache_task = asyncio.create_task(self.update_node_data_cache())
         await self._wait_ready()
 
     async def finalize(self) -> None:

nucliadb/common/cluster/discovery/manual.py

@@ -46,9 +46,7 @@ class ManualDiscovery(AbstractClusterDiscovery):
             except asyncio.CancelledError:
                 return
             except Exception:
-                logger.exception(
-                    "Error while watching cluster members. Will retry at started interval"
-                )
+                logger.exception("Error while watching cluster members. Will retry at started interval")
             finally:
                 await asyncio.sleep(15)
 

nucliadb/common/cluster/discovery/single.py

@@ -48,5 +48,4 @@ class SingleNodeDiscovery(AbstractClusterDiscovery):
             ]
         )
 
-    async def finalize(self) -> None:
-        ...
+    async def finalize(self) -> None: ...

nucliadb/common/cluster/discovery/utils.py

@@ -40,9 +40,7 @@ async def setup_cluster_discovery() -> None:
             # already loaded
             return util
 
-        klass: Union[
-            Type[ManualDiscovery], Type[KubernetesDiscovery], Type[SingleNodeDiscovery]
-        ]
+        klass: Union[Type[ManualDiscovery], Type[KubernetesDiscovery], Type[SingleNodeDiscovery]]
         if settings.cluster_discovery_mode == ClusterDiscoveryMode.MANUAL:
             klass = ManualDiscovery
         elif settings.cluster_discovery_mode == ClusterDiscoveryMode.KUBERNETES:

nucliadb/common/cluster/grpc_node_dummy.py

@@ -22,16 +22,15 @@ from typing import Any
 from nucliadb_protos.nodereader_pb2 import (
     EdgeList,
     RelationEdge,
-    RelationSearchResponse,
 )
-from nucliadb_protos.noderesources_pb2 import EmptyResponse
-from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
 from nucliadb_protos.noderesources_pb2 import (
+    EmptyResponse,
     ShardCreated,
     ShardId,
     ShardIds,
     VectorSetList,
 )
+from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
 from nucliadb_protos.nodewriter_pb2 import OpStatus
 from nucliadb_protos.utils_pb2 import Relation
 
@@ -58,20 +57,23 @@ class DummyWriterStub:  # pragma: no cover
     async def SetResource(self, data):  # pragma: no cover
         self.calls.setdefault("SetResource", []).append(data)
         result = OpStatus()
-        result.field_count = 1
+        return result
+
+    async def SetResourceFromStorage(self, data):  # pragma: no cover
+        self.calls.setdefault("SetResourceFromStorage", []).append(data)
+        result = OpStatus()
         return result
 
     async def AddVectorSet(self, data):  # pragma: no cover
         self.calls.setdefault("AddVectorSet", []).append(data)
         result = OpStatus()
-        result.field_count = 1
         return result
 
-    async def ListVectorSet(self, data: ShardId):  # pragma: no cover
-        self.calls.setdefault("ListVectorSet", []).append(data)
+    async def ListVectorSets(self, data: ShardId):  # pragma: no cover
+        self.calls.setdefault("ListVectorSets", []).append(data)
         result = VectorSetList()
         result.shard.id = data.id
-        result.vectorset.append("base")
+        result.vectorsets.append("base")
         return result
 
     async def GC(self, request: ShardId) -> EmptyResponse:  # pragma: no cover
@@ -87,15 +89,8 @@ class DummyReaderStub:  # pragma: no cover
         self.calls.setdefault("GetShard", []).append(data)
         return NodeResourcesShard(shard_id="shard", fields=2, paragraphs=2, sentences=2)
 
-    async def RelationSearch(self, data):  # pragma: no cover
-        self.calls.setdefault("RelationSearch", []).append(data)
-        result = RelationSearchResponse()
-        return result
-
     async def RelationEdges(self, data):  # pragma: no cover
         self.calls.setdefault("RelationEdges", []).append(data)
         result = EdgeList()
-        result.list.append(
-            RelationEdge(edge_type=Relation.RelationType.ENTITY, property="dummy")
-        )
+        result.list.append(RelationEdge(edge_type=Relation.RelationType.ENTITY, property="dummy"))
         return result

nucliadb/common/cluster/index_node.py

@@ -19,16 +19,13 @@
 #
 from typing import Optional
 
-from lru import LRU  # type: ignore
-from nucliadb_protos.nodereader_pb2_grpc import NodeReaderStub
-from nucliadb_protos.nodewriter_pb2_grpc import NodeWriterStub
+from lru import LRU
 
-from nucliadb.common.cluster.base import AbstractIndexNode  # type: ignore
-from nucliadb.common.cluster.grpc_node_dummy import (  # type: ignore
-    DummyReaderStub,
-    DummyWriterStub,
-)
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.cluster.grpc_node_dummy import DummyReaderStub, DummyWriterStub
 from nucliadb.ingest import SERVICE_NAME
+from nucliadb_protos.nodereader_pb2_grpc import NodeReaderStub
+from nucliadb_protos.nodewriter_pb2_grpc import NodeWriterStub
 from nucliadb_utils.grpc import get_traced_grpc_channel
 
 from .settings import settings
@@ -41,9 +38,7 @@ class IndexNode(AbstractIndexNode):
     _writer: Optional[NodeWriterStub] = None
     _reader: Optional[NodeReaderStub] = None
 
-    def _get_service_address(
-        self, port_map: dict[str, int], port: Optional[int]
-    ) -> str:
+    def _get_service_address(self, port_map: dict[str, int], port: Optional[int]) -> str:
         hostname = self.address.split(":")[0]
         if port is None:
             # For testing purposes we need to be able to have a writing port
@@ -60,10 +55,8 @@ class IndexNode(AbstractIndexNode):
                 grpc_address = self._get_service_address(
                     settings.writer_port_map, settings.node_writer_port
                 )
-                channel = get_traced_grpc_channel(
-                    grpc_address, SERVICE_NAME, variant="_writer"
-                )
-                WRITE_CONNECTIONS[self.address] = NodeWriterStub(channel)  # type: ignore
+                channel = get_traced_grpc_channel(grpc_address, SERVICE_NAME, variant="_writer")
+                WRITE_CONNECTIONS[self.address] = NodeWriterStub(channel)
             else:
                 WRITE_CONNECTIONS[self.address] = DummyWriterStub()
             self._writer = WRITE_CONNECTIONS[self.address]
@@ -76,10 +69,8 @@ class IndexNode(AbstractIndexNode):
                 grpc_address = self._get_service_address(
                     settings.reader_port_map, settings.node_reader_port
                 )
-                channel = get_traced_grpc_channel(
-                    grpc_address, SERVICE_NAME, variant="_reader"
-                )
-                READ_CONNECTIONS[self.address] = NodeReaderStub(channel)  # type: ignore
+                channel = get_traced_grpc_channel(grpc_address, SERVICE_NAME, variant="_reader")
+                READ_CONNECTIONS[self.address] = NodeReaderStub(channel)
             else:
                 READ_CONNECTIONS[self.address] = DummyReaderStub()
             self._reader = READ_CONNECTIONS[self.address]