Skip to content

Add support for managing cluster workload groups #113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 25 commits into from
Jul 17, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
2df97c5
Create model for workload group + policy + request limits policy
ashleyvansp Jul 4, 2025
1daec7b
Add RequestRateLimitPolicies to WorkloadGroup model
ashleyvansp Jul 4, 2025
fdc426a
Add RequestRateLimitsEnforcementPolicy to WorkloadGroup model
ashleyvansp Jul 4, 2025
84f4f3d
Add QueryConsistencyPolicy to the WorkloadGroup model
ashleyvansp Jul 4, 2025
31a3fed
remove unused imports
ashleyvansp Jul 4, 2025
af75cfc
Add create script to workload group
ashleyvansp Jul 4, 2025
865447e
Add workload groups to the Cluster model
ashleyvansp Jul 4, 2025
eb80585
Add deletion script for workload groups
ashleyvansp Jul 4, 2025
e5454dc
Load workload groups from live cluster
ashleyvansp Jul 4, 2025
40fd40c
Setup for deleting workload groups
ashleyvansp Jul 4, 2025
8ba7485
Generate changes for workload groups
ashleyvansp Jul 4, 2025
24555c2
Extract capacity policy changes into its own function
ashleyvansp Jul 6, 2025
044f8e3
Extract workload group changes into its own function
ashleyvansp Jul 6, 2025
ee60410
Add toString methods for workload groups
ashleyvansp Jul 6, 2025
139c088
Update HandleWorkloadGroupChanges
ashleyvansp Jul 6, 2025
c63d44c
Update markdown for a cluster change
ashleyvansp Jul 6, 2025
4ce1148
Add stricter types for RequestRateLimitPolicy properties
ashleyvansp Jul 6, 2025
6a28bb0
Merge branch 'ashleyvansp/workloadGroup' into ashleyvansp/workloadGroups
ashleyvansp Jul 6, 2025
8a2bcb1
Add tests for detecting workload group changes
ashleyvansp Jul 6, 2025
ea09f2b
Add warning to consult support before modifying cluster config
ashleyvansp Jul 6, 2025
8d42dd7
Update README with workload group management
ashleyvansp Jul 6, 2025
df21751
Merge branch 'main' into ashleyvansp/workloadGroups
ashleyvansp Jul 16, 2025
0b44245
Update KustoSchemaTools.Tests/Parser/KustoClusterHandlerTests.cs
ashleyvansp Jul 16, 2025
90b7f7e
Update KustoSchemaTools/Parser/KustoClusterHandler.cs
ashleyvansp Jul 16, 2025
2eeb5c3
Add validations for reading cluster config from YAML
ashleyvansp Jul 16, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
368 changes: 368 additions & 0 deletions KustoSchemaTools.Tests/Changes/ClusterChangesTest.cs
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,315 @@ public void GenerateChanges_WithNullNewCapacityPolicy_ShouldNotGenerateChanges()
Assert.Empty(changeSet.Scripts);
}

#region Workload Group Tests

[Fact]
public void GenerateChanges_WithNewWorkloadGroup_ShouldDetectCreation()
{
// Arrange
var oldCluster = new Cluster { Name = "TestCluster", WorkloadGroups = new List<WorkloadGroup>() };
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 1024)
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.NotEmpty(changeSet.Changes);
Assert.NotEmpty(changeSet.Scripts);

var policyChange = Assert.Single(changeSet.Changes) as PolicyChange<WorkloadGroupPolicy>;
Assert.NotNull(policyChange);
Assert.Equal("test-group", policyChange.Entity);

var scriptContainer = Assert.Single(changeSet.Scripts);
Assert.Contains(".create-or-alter workload_group test-group", scriptContainer.Script.Text);
}

[Fact]
public void GenerateChanges_WithUpdatedWorkloadGroup_ShouldDetectUpdate()
{
// Arrange
var oldCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 1024)
}
};
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 2048)
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.NotEmpty(changeSet.Changes);
Assert.NotEmpty(changeSet.Scripts);

var policyChange = Assert.Single(changeSet.Changes) as PolicyChange<WorkloadGroupPolicy>;
Assert.NotNull(policyChange);
Assert.Equal("test-group", policyChange.Entity);

var scriptContainer = Assert.Single(changeSet.Scripts);
Assert.Contains(".alter-merge workload_group test-group", scriptContainer.Script.Text);
}

[Fact]
public void GenerateChanges_WithIdenticalWorkloadGroups_ShouldDetectNoChanges()
{
// Arrange
var workloadGroup = CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 1024);
var oldCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup> { workloadGroup }
};
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 1024)
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.Empty(changeSet.Changes);
Assert.Empty(changeSet.Scripts);
}

[Fact]
public void GenerateChanges_WithWorkloadGroupDeletion_ShouldDetectDeletion()
{
// Arrange
var oldCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 1024)
}
};
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>(),
Deletions = new ClusterDeletions
{
WorkloadGroups = new List<string> { "test-group" }
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.NotEmpty(changeSet.Changes);

var deletionChange = Assert.Single(changeSet.Changes) as DeletionChange;
Assert.NotNull(deletionChange);
Assert.Equal("test-group", deletionChange.Entity);
Assert.Equal("workload_group", deletionChange.EntityType);
Assert.Contains("Drop test-group", deletionChange.Markdown);
}

[Fact]
public void GenerateChanges_WithWorkloadGroupDeletionOfNonExistentGroup_ShouldNotCreateChange()
{
// Arrange
var oldCluster = new Cluster { Name = "TestCluster", WorkloadGroups = new List<WorkloadGroup>() };
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>(),
Deletions = new ClusterDeletions
{
WorkloadGroups = new List<string> { "non-existent-group" }
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.Empty(changeSet.Changes);
Assert.Empty(changeSet.Scripts);
}

[Fact]
public void GenerateChanges_WithWorkloadGroupMarkedForDeletionButAlsoInNewList_ShouldOnlyProcessDeletion()
{
// Arrange
var oldCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 1024)
}
};
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("test-group", maxMemoryPerQueryPerNode: 2048)
},
Deletions = new ClusterDeletions
{
WorkloadGroups = new List<string> { "test-group" }
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.NotEmpty(changeSet.Changes);

var deletionChange = Assert.Single(changeSet.Changes) as DeletionChange;
Assert.NotNull(deletionChange);
Assert.Equal("test-group", deletionChange.Entity);
Assert.Equal("workload_group", deletionChange.EntityType);
}

[Fact]
public void GenerateChanges_WithMultipleWorkloadGroupChanges_ShouldDetectAllChanges()
{
// Arrange
var oldCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("group1", maxMemoryPerQueryPerNode: 1024),
CreateWorkloadGroup("group2", maxMemoryPerQueryPerNode: 2048),
CreateWorkloadGroup("group3", maxMemoryPerQueryPerNode: 512)
}
};
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateWorkloadGroup("group1", maxMemoryPerQueryPerNode: 1024), // No change
CreateWorkloadGroup("group2", maxMemoryPerQueryPerNode: 4096), // Update
CreateWorkloadGroup("group4", maxMemoryPerQueryPerNode: 256) // New
},
Deletions = new ClusterDeletions
{
WorkloadGroups = new List<string> { "group3" } // Delete
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.Equal(3, changeSet.Changes.Count); // 1 deletion + 1 update + 1 creation

// Check deletion
var deletionChange = changeSet.Changes.OfType<DeletionChange>().Single();
Assert.Equal("group3", deletionChange.Entity);

// Check updates/creations
var policyChanges = changeSet.Changes.OfType<PolicyChange<WorkloadGroupPolicy>>().ToList();
Assert.Equal(2, policyChanges.Count);

var group2Change = policyChanges.First(c => c.Entity == "group2");
var group4Change = policyChanges.First(c => c.Entity == "group4");

Assert.NotNull(group2Change);
Assert.NotNull(group4Change);

// Verify scripts
Assert.Equal(3, changeSet.Scripts.Count);
}

[Fact]
public void GenerateChanges_WithWorkloadGroupHavingComplexPolicy_ShouldDetectChanges()
{
// Arrange
var oldCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateComplexWorkloadGroup("complex-group", 1024, TimeSpan.FromMinutes(5))
}
};
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
CreateComplexWorkloadGroup("complex-group", 2048, TimeSpan.FromMinutes(10))
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.NotEmpty(changeSet.Changes);

var policyChange = Assert.Single(changeSet.Changes) as PolicyChange<WorkloadGroupPolicy>;
Assert.NotNull(policyChange);
Assert.Equal("complex-group", policyChange.Entity);
Assert.Contains("MaxMemoryPerQueryPerNode", policyChange.Markdown);
Assert.Contains("MaxExecutionTime", policyChange.Markdown);
}

[Fact]
public void GenerateChanges_WithWorkloadGroupHavingNullPolicy_ShouldNotCreateChange()
{
// Arrange
var oldCluster = new Cluster { Name = "TestCluster", WorkloadGroups = new List<WorkloadGroup>() };
var newCluster = new Cluster
{
Name = "TestCluster",
WorkloadGroups = new List<WorkloadGroup>
{
new WorkloadGroup { WorkloadGroupName = "null-policy-group", WorkloadGroupPolicy = null }
}
};

// Act
var changeSet = ClusterChanges.GenerateChanges(oldCluster, newCluster, _loggerMock.Object);

// Assert
Assert.NotNull(changeSet);
Assert.Empty(changeSet.Changes);
Assert.Empty(changeSet.Scripts);
}

#endregion

#region Helper Methods
private Cluster CreateClusterWithPolicy(
double? ingestionCapacityCoreUtilizationCoefficient = null,
Expand All @@ -112,6 +421,65 @@ private Cluster CreateClusterWithPolicy(
}
};
}

private WorkloadGroup CreateWorkloadGroup(string name, long? maxMemoryPerQueryPerNode = null)
{
return new WorkloadGroup
{
WorkloadGroupName = name,
WorkloadGroupPolicy = new WorkloadGroupPolicy
{
RequestLimitsPolicy = new RequestLimitsPolicy
{
MaxMemoryPerQueryPerNode = maxMemoryPerQueryPerNode.HasValue
? new PolicyValue<long> { Value = maxMemoryPerQueryPerNode.Value, IsRelaxable = false }
: null
}
}
};
}

private WorkloadGroup CreateComplexWorkloadGroup(string name, long maxMemoryPerQueryPerNode, TimeSpan maxExecutionTime)
{
return new WorkloadGroup
{
WorkloadGroupName = name,
WorkloadGroupPolicy = new WorkloadGroupPolicy
{
RequestLimitsPolicy = new RequestLimitsPolicy
{
MaxMemoryPerQueryPerNode = new PolicyValue<long> { Value = maxMemoryPerQueryPerNode, IsRelaxable = false },
MaxExecutionTime = new PolicyValue<TimeSpan> { Value = maxExecutionTime, IsRelaxable = true },
MaxResultRecords = new PolicyValue<long> { Value = 10000, IsRelaxable = false }
},
RequestRateLimitPolicies = new List<RequestRateLimitPolicy>
{
new RequestRateLimitPolicy
{
IsEnabled = true,
Scope = RateLimitScope.WorkloadGroup,
LimitKind = RateLimitKind.ConcurrentRequests,
Properties = new RateLimitProperties
{
MaxConcurrentRequests = 100
}
},
new RequestRateLimitPolicy
{
IsEnabled = true,
Scope = RateLimitScope.Principal,
LimitKind = RateLimitKind.ResourceUtilization,
Properties = new RateLimitProperties
{
ResourceKind = RateLimitResourceKind.TotalCpuSeconds,
MaxUtilization = 0.8,
TimeWindow = TimeSpan.FromMinutes(5)
}
}
}
}
};
}
#endregion
}
}
Loading
Loading