Add support for managing cluster capacity policy

KustoSchemaTools.Tests/DemoData/DemoDeployment/comprehensive-cluster-policy.yml

@@ -0,0 +1,37 @@
+connections:
+- name: comprehensive-test
+  url: comprehensive.eastus
+capacityPolicy:
+  ingestionCapacity:
+    clusterMaximumConcurrentOperations: 512
+    coreUtilizationCoefficient: 0.75
+  extentsMergeCapacity:
+    minimumConcurrentOperationsPerNode: 1
+    maximumConcurrentOperationsPerNode: 3
+  extentsPurgeRebuildCapacity:
+    maximumConcurrentOperationsPerNode: 1
+  exportCapacity:
+    clusterMaximumConcurrentOperations: 100
+    coreUtilizationCoefficient: 0.25
+  extentsPartitionCapacity:
+    clusterMinimumConcurrentOperations: 1
+    clusterMaximumConcurrentOperations: 32
+  materializedViewsCapacity:
+    clusterMaximumConcurrentOperations: 1
+    extentsRebuildCapacity:
+      clusterMaximumConcurrentOperations: 50
+      maximumConcurrentOperationsPerNode: 5
+  storedQueryResultsCapacity:
+    maximumConcurrentOperationsPerDbAdmin: 250
+    coreUtilizationCoefficient: 0.75
+  streamingIngestionPostProcessingCapacity:
+    maximumConcurrentOperationsPerNode: 4
+  purgeStorageArtifactsCleanupCapacity:
+    maximumConcurrentOperationsPerCluster: 2
+  periodicStorageArtifactsCleanupCapacity:
+    maximumConcurrentOperationsPerCluster: 2
+  queryAccelerationCapacity:
+    clusterMaximumConcurrentOperations: 100
+    coreUtilizationCoefficient: 0.5
+  graphSnapshotsCapacity:
+    clusterMaximumConcurrentOperations: 5

KustoSchemaTools.Tests/DemoData/DemoDeployment/partial-cluster-policy.yml

@@ -0,0 +1,13 @@
+connections:
+- name: partial-test
+  url: partial.eastus
+capacityPolicy:
+  ingestionCapacity:
+    clusterMaximumConcurrentOperations: 256
+  exportCapacity:
+    coreUtilizationCoefficient: 0.5
+  materializedViewsCapacity:
+    clusterMaximumConcurrentOperations: 2
+  queryAccelerationCapacity:
+    clusterMaximumConcurrentOperations: 75
+    coreUtilizationCoefficient: 0.6

KustoSchemaTools/Helpers/Serialization.cs

@@ -14,6 +14,7 @@
         {
             Formatting = Formatting.Indented,
             ContractResolver = new PascalCaseContractResolver(),
+            NullValueHandling = NullValueHandling.Ignore,
         };
 
         public static JsonSerializer JsonSerializer { get; } = new JsonSerializer
@@ -59,19 +60,19 @@
 
         public static T Merge<T>(this T baseObject, T mergeObject)
         {
             var o1 = JObject.FromObject(baseObject, CloneJsonSerializer);
             var o2 = JObject.FromObject(mergeObject, CloneJsonSerializer);
             o1.Merge(o2, new JsonMergeSettings
             {
                 MergeArrayHandling = MergeArrayHandling.Replace 
             });
 
             return o1.ToObject<T>(CloneJsonSerializer);
         }
 
         public static T Clone<T>(this T baseObject)
         {
             return JsonConvert.DeserializeObject<T>(JsonConvert.SerializeObject(baseObject));
         }
 
         public static int RowLength(this string source)

KustoSchemaTools/KustoSchemaHandler.cs

@@ -3,12 +3,10 @@
 using KustoSchemaTools.Model;
 using KustoSchemaTools.Parser;
 using KustoSchemaTools.Parser.KustoLoader;
+using KustoSchemaTools.Parser.KustoWriter;
 using Microsoft.Extensions.Logging;
 using System.Collections.Concurrent;
-using System.Data;
-using System.Diagnostics.Metrics;
 using System.Text;
-using System.Threading.Channels;
 
 namespace KustoSchemaTools
 {
@@ -43,13 +41,13 @@ public KustoSchemaHandler(ILogger<KustoSchemaHandler<T>> schemaHandlerLogger, Ya
             {
                 Log.LogInformation($"Generating diff markdown for {Path.Combine(path, databaseName)} => {cluster}/{escapedDbName}");
 
-                var dbHandler = KustoDatabaseHandlerFactory.Create(cluster.Url, escapedDbName);
+                var dbHandler = CreateDatabaseHandlerWithCapacityPolicy(cluster.Url, escapedDbName, clusters.CapacityPolicy);
                 var kustoDb = await dbHandler.LoadAsync();
                 var changes = DatabaseChanges.GenerateChanges(kustoDb, yamlDb, escapedDbName, Log);
 
                 var comments = changes.Select(itm => itm.Comment).Where(itm => itm != null).ToList();
 
-        
+
                 isValid &= changes.All(itm => itm.Scripts.All(itm => itm.IsValid != false)) && comments.All(itm => itm.FailsRollout == false);
 
                 sb.AppendLine($"# {cluster.Name}/{escapedDbName} ({cluster.Url})");
@@ -74,15 +72,15 @@ public KustoSchemaHandler(ILogger<KustoSchemaHandler<T>> schemaHandlerLogger, Ya
                 }
 
                 var scriptSb = new StringBuilder();
-                foreach(var script in changes.SelectMany(itm => itm.Scripts).Where(itm => itm.IsValid == true).OrderBy(itm => itm.Order))
+                foreach (var script in changes.SelectMany(itm => itm.Scripts).Where(itm => itm.IsValid == true).OrderBy(itm => itm.Order))
                 {
                     scriptSb.AppendLine(script.Text);
                 }
 
                 Log.LogInformation($"Following scripts will be applied:\n{scriptSb}");
             }
 
-            foreach(var follower in yamlDb.Followers)
+            foreach (var follower in yamlDb.Followers)
             {
 
                 Log.LogInformation($"Generating diff markdown for {Path.Combine(path, databaseName)} => {follower.Key}/{follower.Value.DatabaseName}");
@@ -100,7 +98,7 @@ public KustoSchemaHandler(ILogger<KustoSchemaHandler<T>> schemaHandlerLogger, Ya
                 foreach (var change in changes)
                 {
                     sb.AppendLine(change.Markdown);
-                    sb.AppendLine();                    
+                    sb.AppendLine();
                 }
             }
             return (sb.ToString(), isValid);
@@ -112,12 +110,12 @@ public async Task Import(string path, string databaseName, bool includeColumns)
             var clusters = Serialization.YamlPascalCaseDeserializer.Deserialize<Clusters>(clustersFile);
 
             var escapedDbName = databaseName.BracketIfIdentifier();
-            var dbHandler = KustoDatabaseHandlerFactory.Create(clusters.Connections[0].Url, escapedDbName);
+            var dbHandler = CreateDatabaseHandlerWithCapacityPolicy(clusters.Connections[0].Url, escapedDbName, clusters.CapacityPolicy);
 
             var db = await dbHandler.LoadAsync();
             if (includeColumns == false)
             {
-                foreach(var table in db.Tables.Values)
+                foreach (var table in db.Tables.Values)
                 {
                     table.Columns = new Dictionary<string, string>();
                 }
@@ -128,7 +126,7 @@ public async Task Import(string path, string databaseName, bool includeColumns)
         }
 
 
-        public async Task<ConcurrentDictionary<string,Exception>> Apply(string path, string databaseName)
+        public async Task<ConcurrentDictionary<string, Exception>> Apply(string path, string databaseName)
         {
             var clustersFile = File.ReadAllText(Path.Combine(path, "clusters.yml"));
             var clusters = Serialization.YamlPascalCaseDeserializer.Deserialize<Clusters>(clustersFile);
@@ -137,16 +135,18 @@ public async Task<ConcurrentDictionary<string,Exception>> Apply(string path, str
             var yamlHandler = YamlDatabaseHandlerFactory.Create(path, databaseName);
             var yamlDb = await yamlHandler.LoadAsync();
 
-            var results = new ConcurrentDictionary<string,Exception>();
+            var results = new ConcurrentDictionary<string, Exception>();
 
             await Parallel.ForEachAsync(clusters.Connections, async (cluster, token) =>
             {
                 try
                 {
                     Log.LogInformation($"Generating and applying script for {Path.Combine(path, databaseName)} => {cluster}/{escapedDbName}");
-                    var dbHandler = KustoDatabaseHandlerFactory.Create(cluster.Url, escapedDbName);
+
+                    // Create the database handler with capacity policy writer if needed
+                    var dbHandler = CreateDatabaseHandlerWithCapacityPolicy(cluster.Url, escapedDbName, clusters.CapacityPolicy);
                     await dbHandler.WriteAsync(yamlDb);
-                    results.TryAdd(cluster.Url, null);
+                    results.TryAdd(cluster.Url, null!);
                 }
                 catch (Exception ex)
                 {
@@ -156,5 +156,34 @@ await Parallel.ForEachAsync(clusters.Connections, async (cluster, token) =>
 
             return results;
         }
+
+        private IDatabaseHandler<T> CreateDatabaseHandlerWithCapacityPolicy(string clusterUrl, string databaseName, ClusterCapacityPolicy? capacityPolicy)
+        {
+            // Create a logger - in a real application this would come from DI
+            var loggerFactory = LoggerFactory.Create(builder => { });
+            var logger = loggerFactory.CreateLogger<KustoDatabaseHandler<T>>();
+
+            var factory = new KustoDatabaseHandlerFactory<T>(logger)
+                .WithReader<KustoDatabasePrincipalLoader>()
+                .WithReader<KustoDatabaseRetentionAndCacheLoader>()
+                .WithReader<KustoTableBulkLoader>()
+                .WithReader<KustoFunctionBulkLoader>()
+                .WithReader<KustoMaterializedViewBulkLoader>()
+                .WithReader<KustoExternalTableBulkLoader>()
+                .WithReader<KustoContinuousExportBulkLoader>()
+                .WithReader<KustoEntityGroupBulkLoader>()
+                .WithReader<KustoPartitioningPolicyLoader>()
+                .WithReader<ClusterCapacityPolicyLoader>()
+                .WithReader<DatabaseCleanup>()
+                .WithWriter<DefaultDatabaseWriter>();
+
+            // Add capacity policy writer if policy is defined
+            if (capacityPolicy != null)
+            {
+                factory.WithPlugin(new ClusterCapacityPolicyWriter(capacityPolicy));
+            }
+
+            return factory.Create(clusterUrl, databaseName);
+        }
     }
 }

KustoSchemaTools/Model/ClusterCapacityPolicy.cs

@@ -0,0 +1,101 @@
+using KustoSchemaTools.Changes;
+using KustoSchemaTools.Helpers;
+using Newtonsoft.Json;
+
+namespace KustoSchemaTools.Model
+{
+    public class ClusterCapacityPolicy
+    {
+        public IngestionCapacity? IngestionCapacity { get; set; }
+        public ExtentsMergeCapacity? ExtentsMergeCapacity { get; set; }
+        public ExtentsPurgeRebuildCapacity? ExtentsPurgeRebuildCapacity { get; set; }
+        public ExportCapacity? ExportCapacity { get; set; }
+        public ExtentsPartitionCapacity? ExtentsPartitionCapacity { get; set; }
+        public MaterializedViewsCapacity? MaterializedViewsCapacity { get; set; }
+        public StoredQueryResultsCapacity? StoredQueryResultsCapacity { get; set; }
+        public StreamingIngestionPostProcessingCapacity? StreamingIngestionPostProcessingCapacity { get; set; }
+        public PurgeStorageArtifactsCleanupCapacity? PurgeStorageArtifactsCleanupCapacity { get; set; }
+        public PeriodicStorageArtifactsCleanupCapacity? PeriodicStorageArtifactsCleanupCapacity { get; set; }
+        public QueryAccelerationCapacity? QueryAccelerationCapacity { get; set; }
+        public GraphSnapshotsCapacity? GraphSnapshotsCapacity { get; set; }
+
+        public DatabaseScriptContainer CreateScript()
+        {
+            var policy = JsonConvert.SerializeObject(this, Serialization.JsonPascalCase);
+            return new DatabaseScriptContainer("ClusterCapacityPolicy", 10, $".alter-merge cluster policy capacity ```{policy}```");
+        }
+    }
+
+    public class IngestionCapacity
+    {
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+        public double? CoreUtilizationCoefficient { get; set; }
+    }
+
+    public class ExtentsMergeCapacity
+    {
+        public int? MinimumConcurrentOperationsPerNode { get; set; }
+        public int? MaximumConcurrentOperationsPerNode { get; set; }
+    }
+
+    public class ExtentsPurgeRebuildCapacity
+    {
+        public int? MaximumConcurrentOperationsPerNode { get; set; }
+    }
+
+    public class ExportCapacity
+    {
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+        public double? CoreUtilizationCoefficient { get; set; }
+    }
+
+    public class ExtentsPartitionCapacity
+    {
+        public int? ClusterMinimumConcurrentOperations { get; set; }
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+    }
+
+    public class MaterializedViewsCapacity
+    {
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+        public ExtentsRebuildCapacity? ExtentsRebuildCapacity { get; set; }
+    }
+
+    public class ExtentsRebuildCapacity
+    {
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+        public int? MaximumConcurrentOperationsPerNode { get; set; }
+    }
+
+    public class StoredQueryResultsCapacity
+    {
+        public int? MaximumConcurrentOperationsPerDbAdmin { get; set; }
+        public double? CoreUtilizationCoefficient { get; set; }
+    }
+
+    public class StreamingIngestionPostProcessingCapacity
+    {
+        public int? MaximumConcurrentOperationsPerNode { get; set; }
+    }
+
+    public class PurgeStorageArtifactsCleanupCapacity
+    {
+        public int? MaximumConcurrentOperationsPerCluster { get; set; }
+    }
+
+    public class PeriodicStorageArtifactsCleanupCapacity
+    {
+        public int? MaximumConcurrentOperationsPerCluster { get; set; }
+    }
+
+    public class QueryAccelerationCapacity
+    {
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+        public double? CoreUtilizationCoefficient { get; set; }
+    }
+
+    public class GraphSnapshotsCapacity
+    {
+        public int? ClusterMaximumConcurrentOperations { get; set; }
+    }
+}

KustoSchemaTools/Model/Clusters.cs

@@ -4,6 +4,7 @@ public class Clusters
     {
         public List<Cluster> Connections { get; set; } = new List<Cluster>();
         public List<DatabaseScript> Scripts { get; set; } = new List<DatabaseScript>();
+        public ClusterCapacityPolicy? CapacityPolicy { get; set; }
     }
 
 }

KustoSchemaTools/Parser/KustoDatabaseHandlerFactory.cs

@@ -1,4 +1,6 @@
 using KustoSchemaTools.Model;
+using KustoSchemaTools.Parser.KustoLoader;
+using KustoSchemaTools.Parser.KustoWriter;
 using KustoSchemaTools.Plugins;
 using Microsoft.Extensions.Logging;
 
@@ -11,8 +13,8 @@ public KustoDatabaseHandlerFactory(ILogger<KustoDatabaseHandler<T>> logger)
             Logger = logger;
         }
 
-        public List<IKustoBulkEntitiesLoader> Reader { get; } = new ();
-        public List<IDBEntityWriter> Writer { get; } = new ();
+        public List<IKustoBulkEntitiesLoader> Reader { get; } = new();
+        public List<IDBEntityWriter> Writer { get; } = new();
 
         public KustoDatabaseHandlerFactory<T> WithPlugin(IKustoBulkEntitiesLoader plugin)
         {
@@ -43,5 +45,36 @@ public IDatabaseHandler<T> Create(string cluster, string database)
         {
             return new KustoDatabaseHandler<T>(cluster, database, Logger, Reader, Writer);
         }
+
+        public static IDatabaseHandler<T> CreateDefault(string cluster, string database)
+        {
+            // Create a logger - in a real application this would come from DI
+            var loggerFactory = LoggerFactory.Create(builder => { });
+            var logger = loggerFactory.CreateLogger<KustoDatabaseHandler<T>>();
+
+            return new KustoDatabaseHandlerFactory<T>(logger)
+                .WithReader<KustoDatabasePrincipalLoader>()
+                .WithReader<KustoDatabaseRetentionAndCacheLoader>()
+                .WithReader<KustoTableBulkLoader>()
+                .WithReader<KustoFunctionBulkLoader>()
+                .WithReader<KustoMaterializedViewBulkLoader>()
+                .WithReader<KustoExternalTableBulkLoader>()
+                .WithReader<KustoContinuousExportBulkLoader>()
+                .WithReader<KustoEntityGroupBulkLoader>()
+                .WithReader<KustoPartitioningPolicyLoader>()
+                .WithReader<ClusterCapacityPolicyLoader>()
+                .WithReader<DatabaseCleanup>()
+                .WithWriter<DefaultDatabaseWriter>()
+                .Create(cluster, database);
+        }
+    }
+
+    // Add static factory for the non-generic version too
+    public static class KustoDatabaseHandlerFactory
+    {
+        public static IDatabaseHandler<Database> Create(string cluster, string database)
+        {
+            return KustoDatabaseHandlerFactory<Database>.CreateDefault(cluster, database);
+        }
     }
 }

KustoSchemaTools/Parser/KustoLoader/ClusterCapacityPolicyLoader.cs

@@ -0,0 +1,36 @@
+using Kusto.Data.Common;
+using KustoSchemaTools.Model;
+using KustoSchemaTools.Plugins;
+using Newtonsoft.Json;
+
+namespace KustoSchemaTools.Parser.KustoLoader
+{
+    public class ClusterCapacityPolicyLoader : IKustoBulkEntitiesLoader
+    {
+        const string query = @"
+            .show cluster policy capacity
+| project CapacityPolicy = todynamic(Policy)
+        ";
+
+        public async Task Load(Database database, string databaseName, KustoClient client)
+        {
+            try
+            {
+                var response = await client.Client.ExecuteQueryAsync("", query, new ClientRequestProperties());
+                var capacityPolicyData = response.ToScalar<string>();
+
+                if (!string.IsNullOrEmpty(capacityPolicyData))
+                {
+                    var capacityPolicy = JsonConvert.DeserializeObject<ClusterCapacityPolicy>(capacityPolicyData);
+                    // Store in database for comparison if needed
+                    // For now, we don't need to store it as we're only applying changes
+                }
+            }
+            catch (Exception)
+            {
+                // Capacity policy might not exist or access might be denied
+                // This is not critical for the loading process
+            }
+        }
+    }
+}