Skip to content

Commit 7ca6d0b

Browse files
committed
Add shard management functionality and related tests
- Implement methods to get all shards, get a specific shard, and update shard status in CollectionConfigClient. - Add ShardInfo model and ShardStatus enum for shard representation. - Introduce integration tests for shard retrieval and status updates.
1 parent f830ce0 commit 7ca6d0b

File tree

5 files changed

+438
-0
lines changed

5 files changed

+438
-0
lines changed
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
namespace Weaviate.Client.Tests.Integration;
2+
3+
using Weaviate.Client.Models;
4+
5+
public class TestCollectionShards : IntegrationTests
6+
{
7+
private class TestData { }
8+
9+
[Fact]
10+
public async Task Test_Should_Get_Collection_Shards()
11+
{
12+
// Create a collection - it will have at least one shard
13+
var collection = await CollectionFactory<TestData>(
14+
name: "ShardsTest",
15+
properties: [Property.Text("Name")]
16+
);
17+
18+
// Act: Get all shards
19+
var shards = await collection.Config.GetShards(
20+
cancellationToken: TestContext.Current.CancellationToken
21+
);
22+
23+
// Assert
24+
Assert.NotNull(shards);
25+
Assert.NotEmpty(shards);
26+
Assert.All(
27+
shards,
28+
shard =>
29+
{
30+
Assert.NotNull(shard.Name);
31+
Assert.NotEmpty(shard.Name);
32+
Assert.NotNull(shard.Status);
33+
Assert.NotEmpty(shard.Status);
34+
}
35+
);
36+
}
37+
38+
[Fact]
39+
public async Task Test_Should_Get_Specific_Shard()
40+
{
41+
// Arrange: Create collection and get the first shard name
42+
var collection = await CollectionFactory<TestData>(
43+
name: "SpecificShardTest",
44+
properties: [Property.Text("Name")]
45+
);
46+
47+
var allShards = await collection.Config.GetShards(
48+
cancellationToken: TestContext.Current.CancellationToken
49+
);
50+
Assert.NotEmpty(allShards);
51+
52+
var firstShardName = allShards[0].Name;
53+
54+
// Act: Get the specific shard
55+
var shard = await collection.Config.GetShard(
56+
firstShardName,
57+
cancellationToken: TestContext.Current.CancellationToken
58+
);
59+
60+
// Assert
61+
Assert.NotNull(shard);
62+
Assert.Equal(firstShardName, shard.Name);
63+
Assert.NotNull(shard.Status);
64+
Assert.NotEmpty(shard.Status);
65+
}
66+
67+
[Fact]
68+
public async Task Test_Should_Update_Single_Shard_Status_To_ReadOnly()
69+
{
70+
// Arrange: Create collection and get the first shard name
71+
var collection = await CollectionFactory<TestData>(
72+
name: "UpdateShardTest",
73+
properties: [Property.Text("Name")]
74+
);
75+
76+
var allShards = await collection.Config.GetShards(
77+
cancellationToken: TestContext.Current.CancellationToken
78+
);
79+
Assert.NotEmpty(allShards);
80+
81+
var firstShardName = allShards[0].Name;
82+
83+
// Act: Update shard status to READONLY
84+
var updatedShards = await collection.Config.UpdateShardStatus(
85+
ShardStatus.ReadOnly,
86+
firstShardName
87+
);
88+
89+
// Assert
90+
Assert.NotNull(updatedShards);
91+
Assert.Single(updatedShards);
92+
Assert.Equal(firstShardName, updatedShards[0].Name);
93+
Assert.Equal("READONLY", updatedShards[0].Status);
94+
Assert.Equal(ShardStatus.ReadOnly, updatedShards[0].StatusValue);
95+
96+
// Cleanup: Set it back to READY
97+
await collection.Config.UpdateShardStatus(ShardStatus.Ready, firstShardName);
98+
}
99+
100+
[Fact]
101+
public async Task Test_Should_Update_Shard_Status_Back_To_Ready()
102+
{
103+
// Arrange: Create collection, get shard, and set it to READONLY
104+
var collection = await CollectionFactory<TestData>(
105+
name: "UpdateShardBackTest",
106+
properties: [Property.Text("Name")]
107+
);
108+
109+
var allShards = await collection.Config.GetShards(
110+
cancellationToken: TestContext.Current.CancellationToken
111+
);
112+
Assert.NotEmpty(allShards);
113+
114+
var firstShardName = allShards[0].Name;
115+
116+
// Set to READONLY first
117+
await collection.Config.UpdateShardStatus(ShardStatus.ReadOnly, firstShardName);
118+
119+
// Act: Update back to READY
120+
var updatedShards = await collection.Config.UpdateShardStatus(
121+
ShardStatus.Ready,
122+
firstShardName
123+
);
124+
125+
// Assert
126+
Assert.NotNull(updatedShards);
127+
Assert.Single(updatedShards);
128+
Assert.Equal(firstShardName, updatedShards[0].Name);
129+
Assert.Equal("READY", updatedShards[0].Status);
130+
Assert.Equal(ShardStatus.Ready, updatedShards[0].StatusValue);
131+
}
132+
133+
[Fact]
134+
public async Task Test_Should_Update_Multiple_Shards_With_Params()
135+
{
136+
// Arrange: Create collection with multiple shards if possible
137+
var collection = await CollectionFactory<TestData>(
138+
name: "MultipleShardTest",
139+
properties: [Property.Text("Name")],
140+
shardingConfig: new ShardingConfig { DesiredCount = 2 }
141+
);
142+
143+
var allShards = await collection.Config.GetShards(
144+
cancellationToken: TestContext.Current.CancellationToken
145+
);
146+
147+
// If we only have one shard, just test with that one
148+
var shardNames = allShards.Select(s => s.Name).ToArray();
149+
150+
// Act: Update all shards to READONLY
151+
var updatedShards = await collection.Config.UpdateShardStatus(
152+
ShardStatus.ReadOnly,
153+
shardNames
154+
);
155+
156+
// Assert
157+
Assert.NotNull(updatedShards);
158+
Assert.Equal(shardNames.Length, updatedShards.Count);
159+
Assert.All(updatedShards, shard => Assert.Equal("READONLY", shard.Status));
160+
161+
// Cleanup: Set them all back to READY
162+
await collection.Config.UpdateShardStatus(ShardStatus.Ready, shardNames);
163+
}
164+
165+
[Fact]
166+
public async Task Test_Should_Throw_When_No_Shard_Names_Provided()
167+
{
168+
// Arrange
169+
var collection = await CollectionFactory<TestData>(
170+
name: "NoShardNameTest",
171+
properties: [Property.Text("Name")]
172+
);
173+
174+
// Act & Assert
175+
await Assert.ThrowsAsync<ArgumentException>(async () =>
176+
await collection.Config.UpdateShardStatus(ShardStatus.Ready)
177+
);
178+
}
179+
180+
[Fact]
181+
public async Task Test_Should_Throw_When_Shard_Name_Is_Empty()
182+
{
183+
// Arrange
184+
var collection = await CollectionFactory<TestData>(
185+
name: "EmptyShardNameTest",
186+
properties: [Property.Text("Name")]
187+
);
188+
189+
// Act & Assert
190+
await Assert.ThrowsAsync<ArgumentException>(async () =>
191+
await collection.Config.GetShard(
192+
"",
193+
cancellationToken: TestContext.Current.CancellationToken
194+
)
195+
);
196+
}
197+
}

src/Weaviate.Client/Models/Collection.Update.cs

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,4 +431,73 @@ await _client.Collections.Export(_collectionName)
431431

432432
return response?.ToModel();
433433
}
434+
435+
/// <summary>
436+
/// Gets all shards for this collection.
437+
/// </summary>
438+
/// <param name="cancellationToken">Cancellation token to cancel the operation.</param>
439+
/// <returns>A list of shard information for the collection.</returns>
440+
public async Task<IList<ShardInfo>> GetShards(CancellationToken cancellationToken = default)
441+
{
442+
var shards = await _client.RestClient.CollectionGetShards(
443+
_collectionName,
444+
cancellationToken
445+
);
446+
447+
return shards.Select(s => s.ToModel()).ToList();
448+
}
449+
450+
/// <summary>
451+
/// Gets information about a specific shard.
452+
/// </summary>
453+
/// <param name="shardName">The name of the shard to retrieve.</param>
454+
/// <param name="cancellationToken">Cancellation token to cancel the operation.</param>
455+
/// <returns>Information about the specified shard, or null if not found.</returns>
456+
public async Task<ShardInfo?> GetShard(
457+
string shardName,
458+
CancellationToken cancellationToken = default
459+
)
460+
{
461+
ArgumentException.ThrowIfNullOrEmpty(shardName);
462+
463+
var shard = await _client.RestClient.CollectionGetShard(
464+
_collectionName,
465+
shardName,
466+
cancellationToken
467+
);
468+
469+
return shard?.ToModel();
470+
}
471+
472+
/// <summary>
473+
/// Updates the status of one or more shards.
474+
/// </summary>
475+
/// <param name="status">The new status to set for the shards.</param>
476+
/// <param name="shardNames">The names of the shards to update.</param>
477+
/// <returns>A list of updated shard information.</returns>
478+
public async Task<IList<ShardInfo>> UpdateShardStatus(
479+
ShardStatus status,
480+
params string[] shardNames
481+
)
482+
{
483+
ArgumentNullException.ThrowIfNull(shardNames);
484+
485+
if (shardNames.Length == 0)
486+
{
487+
throw new ArgumentException(
488+
"At least one shard name must be provided.",
489+
nameof(shardNames)
490+
);
491+
}
492+
493+
var shardStatus = new Rest.Dto.ShardStatus { Status = status.ToApiString() };
494+
495+
var tasks = shardNames.Select(shardName =>
496+
_client.RestClient.CollectionUpdateShard(_collectionName, shardName, shardStatus)
497+
);
498+
499+
var results = await Task.WhenAll(tasks);
500+
501+
return results.Select(r => r.ToModel()).ToList();
502+
}
434503
}
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
namespace Weaviate.Client.Models;
2+
3+
/// <summary>
4+
/// Represents the possible status values for a shard.
5+
/// </summary>
6+
public enum ShardStatus
7+
{
8+
/// <summary>
9+
/// The shard is ready to serve requests.
10+
/// </summary>
11+
Ready,
12+
13+
/// <summary>
14+
/// The shard is in read-only mode.
15+
/// </summary>
16+
ReadOnly,
17+
18+
/// <summary>
19+
/// The shard is indexing data.
20+
/// </summary>
21+
Indexing,
22+
}
23+
24+
/// <summary>
25+
/// Information about a collection shard, including its name, status, and statistics.
26+
/// </summary>
27+
public record ShardInfo
28+
{
29+
/// <summary>
30+
/// Gets or sets the name of the shard.
31+
/// </summary>
32+
public string Name { get; set; } = string.Empty;
33+
34+
/// <summary>
35+
/// Gets or sets the status of the shard (e.g., "READY", "READONLY", "INDEXING").
36+
/// </summary>
37+
public string Status { get; set; } = string.Empty;
38+
39+
/// <summary>
40+
/// Gets or sets the size of the vector queue for the shard.
41+
/// </summary>
42+
public int? VectorQueueSize { get; set; }
43+
44+
/// <summary>
45+
/// Gets the parsed status as a ShardStatus enum.
46+
/// Returns null if the status cannot be parsed.
47+
/// </summary>
48+
public ShardStatus? StatusValue =>
49+
Status?.ToUpperInvariant() switch
50+
{
51+
"READY" => ShardStatus.Ready,
52+
"READONLY" => ShardStatus.ReadOnly,
53+
"INDEXING" => ShardStatus.Indexing,
54+
_ => null,
55+
};
56+
}
57+
58+
/// <summary>
59+
/// Extension methods for ShardStatus enum.
60+
/// </summary>
61+
internal static class ShardStatusExtensions
62+
{
63+
/// <summary>
64+
/// Converts a ShardStatus enum to its string representation for the API.
65+
/// </summary>
66+
internal static string ToApiString(this ShardStatus status)
67+
{
68+
return status switch
69+
{
70+
ShardStatus.Ready => "READY",
71+
ShardStatus.ReadOnly => "READONLY",
72+
ShardStatus.Indexing => "INDEXING",
73+
_ => throw new ArgumentOutOfRangeException(
74+
nameof(status),
75+
status,
76+
"Unknown shard status value"
77+
),
78+
};
79+
}
80+
}

0 commit comments

Comments
 (0)