def test_parse_shards_valid2(self) -> None:
# Parse multiple inputs
self.assertEqual(
{
ShardID(node=NodeID(node_index=0), shard_index=1),
ShardID(node=NodeID(node_index=1), shard_index=2),
},
helpers.parse_shards(["N0:S1", "N1:S2"]),
)
# Remove duplicates
self.assertEqual(
{
ShardID(node=NodeID(node_index=0), shard_index=1),
ShardID(node=NodeID(node_index=1), shard_index=2),
},
helpers.parse_shards(["N0:S1", "N1:S2", "N0:s1"]),
)
def _combine(
cv: ClusterView,
shards: Optional[List[str]] = None,
node_names: Optional[List[str]] = None,
node_indexes: Optional[List[int]] = None,
) -> Tuple[ShardID, ...]:
shards = list(shards or [])
node_names = list(node_names or [])
node_indexes = list(node_indexes or [])
shard_ids = parse_shards(shards)
for nn in node_names:
shard_ids.add(ShardID(node=cv.get_node_id(node_name=nn), shard_index=-1))
for ni in node_indexes:
shard_ids.add(ShardID(node=NodeID(node_index=ni), shard_index=-1))
shard_ids_expanded = cv.expand_shards(shard_ids)
return shard_ids_expanded
def test_parse_shards_invalid(self) -> None:
"""
Invalid shards should throw ValueError exceptions.
"""
# N0:S1, N0, or 0 format should be accepted
with self.assertRaises(ValueError):
print(helpers.parse_shards([":S2"]))
with self.assertRaises(ValueError):
print(helpers.parse_shards(["N:S2"]))
with self.assertRaises(ValueError):
print(helpers.parse_shards(["X0"]))
with self.assertRaises(ValueError):
helpers.parse_shards(["N0:B1"])
with self.assertRaises(ValueError):
helpers.parse_shards(["N0:S1X"])
def test_parse_shards_valid1(self) -> None:
# 5
self.assertEqual(
{ShardID(node=NodeID(node_index=5), shard_index=ALL_SHARDS)},
helpers.parse_shards(["5"]),
)
# 5:1
self.assertEqual(
{ShardID(node=NodeID(node_index=5), shard_index=1)},
helpers.parse_shards(["5:1"]),
)
# 0:S1
self.assertEqual(
{ShardID(node=NodeID(node_index=0), shard_index=1)},
helpers.parse_shards(["0:S1"]),
)
# N0:S1
self.assertEqual(
{ShardID(node=NodeID(node_index=0), shard_index=1)},
helpers.parse_shards(["N0:S1"]),
)
# N0 == ShardID(0, ALL_SHARDS)
self.assertEqual(
{ShardID(node=NodeID(node_index=0), shard_index=ALL_SHARDS)},
helpers.parse_shards(["N0"]),
)
# N1:S4 == ShardID(1, 4)
self.assertEqual(
{ShardID(node=NodeID(node_index=1), shard_index=4)},
helpers.parse_shards(["N1:S4"]),
)
# Allow ignored case
# n1:S4 == ShardID(1, 4)
self.assertEqual(
{ShardID(node=NodeID(node_index=1), shard_index=4)},
helpers.parse_shards(["n1:S4"]),
)
async def apply(
self,
reason: str,
node_indexes: Optional[List[int]] = None,
node_names: Optional[List[str]] = None,
shards: Optional[List[str]] = None,
shard_target_state: Optional[str] = "may-disappear",
sequencer_node_indexes: Optional[List[int]] = None,
sequencer_node_names: Optional[List[str]] = None,
user: Optional[str] = "",
group: Optional[bool] = True,
skip_safety_checks: Optional[bool] = False,
ttl: Optional[int] = 0,
allow_passive_drains: Optional[bool] = False,
force_restore_rebuilding: Optional[bool] = False,
):
"""
Applies new maintenance to Maintenance Manager
"""
ctx = context.get_context()
try:
async with ctx.get_cluster_admin_client() as client:
cv = await get_cluster_view(client)
all_node_indexes = set()
if node_indexes is not None:
all_node_indexes = all_node_indexes.union(set(node_indexes))
if node_names is not None:
all_node_indexes = all_node_indexes.union({
cv.get_node_index(node_name=node_name)
for node_name in set(node_names)
})
shard_ids = set()
sequencer_nodes = set()
for ni in all_node_indexes:
nv = cv.get_node_view(node_index=ni)
if nv.is_storage:
shard_ids.add(ShardID(node=nv.node_id, shard_index=-1))
if nv.is_sequencer:
sequencer_nodes.add(nv.node_id)
if sequencer_node_indexes is not None:
for ni in set(sequencer_node_indexes):
nv = cv.get_node_view(node_index=ni)
if nv.is_sequencer:
sequencer_nodes.add(nv.node_id)
if sequencer_node_names is not None:
for nn in set(sequencer_node_names):
nv = cv.get_node_view(node_name=nn)
if nv.is_sequencer:
sequencer_nodes.add(nv.node_id)
if shards is not None:
shard_ids = shard_ids.union(
cv.expand_shards(parse_shards(shards)))
except NodeNotFoundError as e:
print(colored(f"Node not found: {e}", "red"))
return
try:
async with ctx.get_cluster_admin_client() as client:
maintenances: Collection[MaintenanceDefinition]
maintenances = await apply_maintenance(
client=client,
shards=shard_ids,
shard_target_state=_parse_shard_target_state(
shard_target_state),
sequencer_nodes=list(sequencer_nodes),
group=group,
ttl=timedelta(seconds=ttl),
user=user or getuser(),
reason=reason,
skip_safety_checks=skip_safety_checks,
allow_passive_drains=allow_passive_drains,
force_restore_rebuilding=force_restore_rebuilding,
)
cv = await get_cluster_view(client)
except Exception as e:
print(colored(f"Cannot apply maintenance: {e}", "red"))
return
print(
_render(
[
cv.get_maintenance_view_by_id(id)
for id in [mnt.group_id for mnt in maintenances]
],
cv,
mode=RenderingMode.EXPANDED,
))
def test_parse_shards_empty(self) -> None:
"""
Empty input should produce empty output
"""
self.assertEqual(set(), helpers.parse_shards([]))