Add shard management functionality

src/Weaviate.Client.Tests/Integration/TestCollectionShards.cs

@@ -0,0 +1,194 @@
+namespace Weaviate.Client.Tests.Integration;
+
+using Weaviate.Client.Models;
+
+public class TestCollectionShards : IntegrationTests
+{
+    private class TestData { }
+
+    [Fact]
+    public async Task Test_Should_Get_Collection_Shards()
+    {
+        // Create a collection - it will have at least one shard
+        var collection = await CollectionFactory<TestData>(
+            name: "ShardsTest",
+            properties: [Property.Text("Name")]
+        );
+
+        // Act: Get all shards
+        var shards = await collection.Config.GetShards(
+            cancellationToken: TestContext.Current.CancellationToken
+        );
+
+        // Assert
+        Assert.NotNull(shards);
+        Assert.NotEmpty(shards);
+        Assert.All(
+            shards,
+            shard =>
+            {
+                Assert.NotNull(shard.Name);
+                Assert.NotEmpty(shard.Name);
+                // Status is an enum, so it always has a value
+                Assert.True(Enum.IsDefined(typeof(ShardStatus), shard.Status));
+            }
+        );
+    }
+
+    [Fact]
+    public async Task Test_Should_Get_Specific_Shard()
+    {
+        // Arrange: Create collection and get the first shard name
+        var collection = await CollectionFactory<TestData>(
+            name: "SpecificShardTest",
+            properties: [Property.Text("Name")]
+        );
+
+        var allShards = await collection.Config.GetShards(
+            cancellationToken: TestContext.Current.CancellationToken
+        );
+        Assert.NotEmpty(allShards);
+
+        var firstShardName = allShards[0].Name;
+
+        // Act: Get the specific shard
+        var shard = await collection.Config.GetShard(
+            firstShardName,
+            cancellationToken: TestContext.Current.CancellationToken
+        );
+
+        // Assert
+        Assert.NotNull(shard);
+        Assert.Equal(firstShardName, shard.Name);
+        Assert.True(Enum.IsDefined(typeof(ShardStatus), shard.Status));
+    }
+
+    [Fact]
+    public async Task Test_Should_Update_Single_Shard_Status_To_ReadOnly()
+    {
+        // Arrange: Create collection and get the first shard name
+        var collection = await CollectionFactory<TestData>(
+            name: "UpdateShardTest",
+            properties: [Property.Text("Name")]
+        );
+
+        var allShards = await collection.Config.GetShards(
+            cancellationToken: TestContext.Current.CancellationToken
+        );
+        Assert.NotEmpty(allShards);
+
+        var firstShardName = allShards[0].Name;
+
+        // Act: Update shard status to READONLY
+        var updatedShards = await collection.Config.UpdateShardStatus(
+            ShardStatus.ReadOnly,
+            firstShardName
+        );
+
+        // Assert
+        Assert.NotNull(updatedShards);
+        Assert.Single(updatedShards);
+        Assert.Equal(firstShardName, updatedShards[0].Name);
+        Assert.Equal(ShardStatus.ReadOnly, updatedShards[0].Status);
+
+        // Cleanup: Set it back to READY
+        await collection.Config.UpdateShardStatus(ShardStatus.Ready, firstShardName);
+    }
+
+    [Fact]
+    public async Task Test_Should_Update_Shard_Status_Back_To_Ready()
+    {
+        // Arrange: Create collection, get shard, and set it to READONLY
+        var collection = await CollectionFactory<TestData>(
+            name: "UpdateShardBackTest",
+            properties: [Property.Text("Name")]
+        );
+
+        var allShards = await collection.Config.GetShards(
+            cancellationToken: TestContext.Current.CancellationToken
+        );
+        Assert.NotEmpty(allShards);
+
+        var firstShardName = allShards[0].Name;
+
+        // Set to READONLY first
+        await collection.Config.UpdateShardStatus(ShardStatus.ReadOnly, firstShardName);
+
+        // Act: Update back to READY
+        var updatedShards = await collection.Config.UpdateShardStatus(
+            ShardStatus.Ready,
+            firstShardName
+        );
+
+        // Assert
+        Assert.NotNull(updatedShards);
+        Assert.Single(updatedShards);
+        Assert.Equal(firstShardName, updatedShards[0].Name);
+        Assert.Equal(ShardStatus.Ready, updatedShards[0].Status);
+    }
+
+    [Fact]
+    public async Task Test_Should_Update_Multiple_Shards_With_Params()
+    {
+        // Arrange: Create collection with multiple shards if possible
+        var collection = await CollectionFactory<TestData>(
+            name: "MultipleShardTest",
+            properties: [Property.Text("Name")],
+            shardingConfig: new ShardingConfig { DesiredCount = 2 }
+        );
+
+        var allShards = await collection.Config.GetShards(
+            cancellationToken: TestContext.Current.CancellationToken
+        );
+
+        // If we only have one shard, just test with that one
+        var shardNames = allShards.Select(s => s.Name).ToArray();
+
+        // Act: Update all shards to READONLY
+        var updatedShards = await collection.Config.UpdateShardStatus(
+            ShardStatus.ReadOnly,
+            shardNames
+        );
+
+        // Assert
+        Assert.NotNull(updatedShards);
+        Assert.Equal(shardNames.Length, updatedShards.Count);
+        Assert.All(updatedShards, shard => Assert.Equal(ShardStatus.ReadOnly, shard.Status));
+
+        // Cleanup: Set them all back to READY
+        await collection.Config.UpdateShardStatus(ShardStatus.Ready, shardNames);
+    }
+
+    [Fact]
+    public async Task Test_Should_Throw_When_No_Shard_Names_Provided()
+    {
+        // Arrange
+        var collection = await CollectionFactory<TestData>(
+            name: "NoShardNameTest",
+            properties: [Property.Text("Name")]
+        );
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentException>(async () =>
+            await collection.Config.UpdateShardStatus(ShardStatus.Ready)
+        );
+    }
+
+    [Fact]
+    public async Task Test_Should_Throw_When_Shard_Name_Is_Empty()
+    {
+        // Arrange
+        var collection = await CollectionFactory<TestData>(
+            name: "EmptyShardNameTest",
+            properties: [Property.Text("Name")]
+        );
+
+        // Act & Assert
+        await Assert.ThrowsAsync<ArgumentException>(async () =>
+            await collection.Config.GetShard(
+                "",
+                cancellationToken: TestContext.Current.CancellationToken
+            )
+        );
+    }
+}

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

@@ -408,4 +408,73 @@ await _client.Collections.Export(_collectionName)
 
         return response?.ToModel();
     }
+
+    /// <summary>
+    /// Gets all shards for this collection.
+    /// </summary>
+    /// <param name="cancellationToken">Cancellation token to cancel the operation.</param>
+    /// <returns>A list of shard information for the collection.</returns>
+    public async Task<IList<ShardInfo>> GetShards(CancellationToken cancellationToken = default)
+    {
+        var shards = await _client.RestClient.CollectionGetShards(
+            _collectionName,
+            cancellationToken
+        );
+
+        return shards.Select(s => s.ToModel()).ToList();
+    }
+
+    /// <summary>
+    /// Gets information about a specific shard.
+    /// </summary>
+    /// <param name="shardName">The name of the shard to retrieve.</param>
+    /// <param name="cancellationToken">Cancellation token to cancel the operation.</param>
+    /// <returns>Information about the specified shard, or null if not found.</returns>
+    public async Task<ShardInfo?> GetShard(
+        string shardName,
+        CancellationToken cancellationToken = default
+    )
+    {
+        ArgumentException.ThrowIfNullOrEmpty(shardName);
+
+        var shard = await _client.RestClient.CollectionGetShard(
+            _collectionName,
+            shardName,
+            cancellationToken
+        );
+
+        return shard?.ToModel();
+    }
+
+    /// <summary>
+    /// Updates the status of one or more shards.
+    /// </summary>
+    /// <param name="status">The new status to set for the shards.</param>
+    /// <param name="shardNames">The names of the shards to update.</param>
+    /// <returns>A list of updated shard information.</returns>
+    public async Task<IList<ShardInfo>> UpdateShardStatus(
+        ShardStatus status,
+        params string[] shardNames
+    )
+    {
+        ArgumentNullException.ThrowIfNull(shardNames);
+
+        if (shardNames.Length == 0)
+        {
+            throw new ArgumentException(
+                "At least one shard name must be provided.",
+                nameof(shardNames)
+            );
+        }
+
+        var shardStatus = new Rest.Dto.ShardStatus { Status = status.ToApiString() };
+
+        var tasks = shardNames.Select(shardName =>
+            _client.RestClient.CollectionUpdateShard(_collectionName, shardName, shardStatus)
+        );
+
+        var results = await Task.WhenAll(tasks);
+
+        return results.Select(r => r.ToModel()).ToList();
+    }
 }

src/Weaviate.Client/Models/Shard.cs

@@ -0,0 +1,68 @@
+using System.Runtime.Serialization;
+
+namespace Weaviate.Client.Models;
+
+/// <summary>
+/// Represents the possible status values for a shard.
+/// </summary>
+public enum ShardStatus
+{
+    /// <summary>
+    /// The shard is ready to serve requests.
+    /// </summary>
+    [EnumMember(Value = "READY")]
+    Ready,
+
+    /// <summary>
+    /// The shard is in read-only mode.
+    /// </summary>
+    [EnumMember(Value = "READONLY")]
+    ReadOnly,
+
+    /// <summary>
+    /// The shard is indexing data.
+    /// </summary>
+    [EnumMember(Value = "INDEXING")]
+    Indexing,
+}
+
+/// <summary>
+/// Information about a collection shard, including its name, status, and statistics.
+/// </summary>
+public record ShardInfo
+{
+    /// <summary>
+    /// Gets or sets the name of the shard.
+    /// </summary>
+    public string Name { get; set; } = string.Empty;
+
+    /// <summary>
+    /// Gets or sets the status of the shard.
+    /// </summary>
+    public ShardStatus Status { get; set; } = ShardStatus.Ready;
+
+    /// <summary>
+    /// Gets or sets the size of the vector queue for the shard.
+    /// </summary>
+    public int? VectorQueueSize { get; set; }
+}
+
+/// <summary>
+/// Extension methods for ShardStatus enum.
+/// </summary>
+internal static class ShardStatusExtensions
+{
+    /// <summary>
+    /// Converts a ShardStatus enum to its string representation for the API.
+    /// </summary>
+    internal static string ToApiString(this ShardStatus status) =>
+        WeaviateExtensions.ToEnumMemberString(status);
+
+    /// <summary>
+    /// Parses a string value to a ShardStatus enum.
+    /// </summary>
+    internal static ShardStatus ParseStatus(string? value) =>
+        string.IsNullOrEmpty(value)
+            ? ShardStatus.Ready
+            : WeaviateExtensions.FromEnumMemberString<ShardStatus>(value);
+}