From ed080cf8e36866bb83eb24772f317639ae2f378c Mon Sep 17 00:00:00 2001 From: Ferran Llamas Date: Fri, 12 Jan 2024 15:07:22 +0100 Subject: [PATCH] Return secondary nodes too on cluster members endpoint (#1723) --- nucliadb/nucliadb/common/cluster/base.py | 5 +- nucliadb/nucliadb/common/cluster/manager.py | 7 ++- nucliadb/nucliadb/search/app.py | 3 +- .../nucliadb/search/tests/unit/test_app.py | 28 +++++++++-- .../tests/unit/common/cluster/test_manager.py | 46 +++++++++++++++++++ 5 files changed, 81 insertions(+), 8 deletions(-) diff --git a/nucliadb/nucliadb/common/cluster/base.py b/nucliadb/nucliadb/common/cluster/base.py index 048722260d..3825deac57 100644 --- a/nucliadb/nucliadb/common/cluster/base.py +++ b/nucliadb/nucliadb/common/cluster/base.py @@ -50,7 +50,10 @@ def __init__( self.primary_id = primary_id def __str__(self): - return f"{self.__class__.__name__}({self.id}, {self.address})" + if self.primary_id is None: + return f"{self.__class__.__name__}({self.id}, {self.address})" + else: + return f"{self.__class__.__name__}({self.id}, {self.address}, primary_id={self.primary_id})" def __repr__(self): return self.__str__() diff --git a/nucliadb/nucliadb/common/cluster/manager.py b/nucliadb/nucliadb/common/cluster/manager.py index 7a27ebc176..f2e65275fa 100644 --- a/nucliadb/nucliadb/common/cluster/manager.py +++ b/nucliadb/nucliadb/common/cluster/manager.py @@ -63,8 +63,11 @@ READ_REPLICA_INDEX_NODES: dict[str, set[str]] = {} -def get_index_nodes() -> list[AbstractIndexNode]: - return [inode for inode in INDEX_NODES.values() if inode.primary_id is None] +def get_index_nodes(include_secondary: bool = False) -> list[AbstractIndexNode]: + all_nodes = [inode for inode in INDEX_NODES.values()] + if not include_secondary: + return [inode for inode in all_nodes if inode.primary_id is None] + return all_nodes def get_index_node(node_id: str) -> Optional[AbstractIndexNode]: diff --git a/nucliadb/nucliadb/search/app.py b/nucliadb/nucliadb/search/app.py index f19102c6d0..e69a891071 100644 --- a/nucliadb/nucliadb/search/app.py +++ b/nucliadb/nucliadb/search/app.py @@ -116,8 +116,9 @@ async def node_members(request: Request) -> JSONResponse: "type": node.label, "shard_count": node.shard_count, "dummy": node.dummy, + "primary_id": node.primary_id, } - for node in manager.get_index_nodes() + for node in manager.get_index_nodes(include_secondary=True) ] ) diff --git a/nucliadb/nucliadb/search/tests/unit/test_app.py b/nucliadb/nucliadb/search/tests/unit/test_app.py index f87249ee9f..a2bdaa95dc 100644 --- a/nucliadb/nucliadb/search/tests/unit/test_app.py +++ b/nucliadb/nucliadb/search/tests/unit/test_app.py @@ -17,34 +17,54 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . +import json from unittest.mock import patch import pytest +from nucliadb.common.cluster.index_node import IndexNode from nucliadb.search import app pytestmark = pytest.mark.asyncio async def test_alive(): - with patch.object(app.manager, "get_index_nodes", return_value={"node1": "node1"}): + with patch.object(app.manager, "get_index_nodes", return_value=[{"id": "node1"}]): resp = await app.alive(None) assert resp.status_code == 200 async def test_not_alive(): - with patch.object(app.manager, "get_index_nodes", return_value={}): + with patch.object(app.manager, "get_index_nodes", return_value=[]): resp = await app.alive(None) assert resp.status_code == 503 async def test_ready(): - with patch.object(app.manager, "get_index_nodes", return_value={"node1": "node1"}): + with patch.object(app.manager, "get_index_nodes", return_value=[{"id": "node1"}]): resp = await app.ready(None) assert resp.status_code == 200 async def test_not_ready(): - with patch.object(app.manager, "get_index_nodes", return_value={}): + with patch.object(app.manager, "get_index_nodes", return_value=[]): resp = await app.ready(None) assert resp.status_code == 503 + + +async def test_node_members(): + nodes = [ + IndexNode(id="node1", address="node1", shard_count=0, dummy=True), + IndexNode( + id="node2", address="node2", shard_count=0, dummy=True, primary_id="node1" + ), + ] + with patch.object(app.manager, "get_index_nodes", return_value=nodes): + resp = await app.node_members(None) + assert resp.status_code == 200 + members = json.loads(resp.body) + sorted(members, key=lambda x: x["id"]) + assert members[0]["id"] == "node1" + assert members[0]["primary_id"] is None + assert members[1]["id"] == "node2" + assert members[1]["primary_id"] == "node1" diff --git a/nucliadb/nucliadb/tests/unit/common/cluster/test_manager.py b/nucliadb/nucliadb/tests/unit/common/cluster/test_manager.py index 08a29afcd5..ab3b8ce156 100644 --- a/nucliadb/nucliadb/tests/unit/common/cluster/test_manager.py +++ b/nucliadb/nucliadb/tests/unit/common/cluster/test_manager.py @@ -18,6 +18,7 @@ # along with this program. If not, see . # import asyncio +from unittest import mock from unittest.mock import MagicMock import pytest @@ -174,3 +175,48 @@ def test_choose_node(): ) manager.INDEX_NODES.clear() + + +@pytest.fixture(scope="function") +def standalone_mode_off(): + prev = settings.standalone_mode + settings.standalone_mode = False + yield + settings.standalone_mode = prev + + +@pytest.fixture(scope="function") +def index_nodes(): + index_nodes = {} + with mock.patch.object(manager, "INDEX_NODES", new=index_nodes): + yield index_nodes + + +def test_get_index_nodes(standalone_mode_off, index_nodes): + # Add a primary node + manager.add_index_node( + id="node-0", + address="nohost", + shard_count=0, + dummy=True, + ) + # Add a secondary replica of node-0 + manager.add_index_node( + id="node-1", + address="nohost", + shard_count=0, + dummy=True, + primary_id="node-0", + ) + + # By default, only primary nodes are returned + nodes = manager.get_index_nodes() + assert len(nodes) == 1 + assert nodes[0].id == "node-0" + + # If we ask for secondary, we get both + nodes = manager.get_index_nodes(include_secondary=True) + assert len(nodes) == 2 + sorted(nodes, key=lambda x: x.id) + assert nodes[0].id == "node-0" + assert nodes[1].id == "node-1"