Add schema discovery command and read-only MCP tool (#160)

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### New features
 
+- **`schema` discovery command and MCP tool.** A new read-only `schema` command infers a container's structure from a small, bounded sample: it returns the partition key path(s), an indexing policy summary, an estimated document count, and inferred field types (with dot notation for nested objects and per-field presence counts). `--sample <n>` selects how many documents to sample (clamped to 1-100, default 20), and `--database`/`--container` override the target. Exposed as a read-only MCP tool so agents can discover container structure cheaply instead of re-sampling or guessing field names. ([#160](https://github.com/Azure/CosmosDBShell/issues/160))
 - `bucket` command now manages container throughput bucket limits in addition to client-side bucket selection. `bucket show` lists the throughput bucket limits configured on the current container, `bucket set <1-5> <1-100>` limits a bucket to a maximum percentage of the container's throughput, and `bucket clear <1-5>` removes a bucket's limit. These control-plane subcommands target the current container (or `--container`) and require an Azure AD (Entra) connection; the existing client-side `bucket`, `bucket <1-5>`, and `bucket 0` selection continues to work on any connection. ([#144](https://github.com/Azure/CosmosDBShell/issues/144))
 - `throughput` write subcommands (`set`/`manual`/`autoscale`) now accept `--dry-run` to preview the change — reporting current vs. planned mode and RU/s as JSON (and a table interactively) — without applying it or prompting for confirmation. A first slice of dry-run mode (item G1). ([#164](https://github.com/Azure/CosmosDBShell/issues/164))
 - **`--dry-run` for destructive delete commands.** `rm`, `rmcon`, `rmdb`, and `delete` accept `--dry-run` to preview the effect without deleting anything: `rm` reports how many items match the pattern, and `rmcon`/`rmdb` report the container or database that would be removed. No confirmation prompt is shown and no changes are made. ([#156](https://github.com/Azure/CosmosDBShell/issues/156))

CosmosDBShell.Tests/CommandTests/SchemaCommandTests.cs

@@ -0,0 +1,153 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace CosmosShell.Tests.CommandTests;
+
+using System.Text.Json;
+using Azure.Data.Cosmos.Shell.Commands;
+using Azure.Data.Cosmos.Shell.Core;
+
+/// <summary>
+/// Unit tests for <see cref="SchemaCommand"/>. Covers the pure helpers that clamp the
+/// sample size and infer field types, which can be exercised without a live Cosmos DB
+/// connection.
+/// </summary>
+public class SchemaCommandTests
+{
+    [Fact]
+    public void SchemaCommand_IsRegistered()
+    {
+        var runner = new CommandRunner();
+
+        Assert.True(runner.Commands.TryGetValue("schema", out var factory));
+        Assert.Equal("schema", factory!.CommandName);
+    }
+
+    [Fact]
+    public void NormalizeSample_UsesDefaultWhenMissing()
+    {
+        Assert.Equal(SchemaCommand.DefaultSample, SchemaCommand.NormalizeSample(null));
+    }
+
+    [Fact]
+    public void NormalizeSample_ClampsBelowMinimum()
+    {
+        Assert.Equal(SchemaCommand.MinSample, SchemaCommand.NormalizeSample(0));
+        Assert.Equal(SchemaCommand.MinSample, SchemaCommand.NormalizeSample(-5));
+    }
+
+    [Fact]
+    public void NormalizeSample_ClampsAboveMaximum()
+    {
+        Assert.Equal(SchemaCommand.MaxSample, SchemaCommand.NormalizeSample(SchemaCommand.MaxSample + 1));
+        Assert.Equal(SchemaCommand.MaxSample, SchemaCommand.NormalizeSample(int.MaxValue));
+    }
+
+    [Fact]
+    public void NormalizeSample_KeepsValueInRange()
+    {
+        Assert.Equal(42, SchemaCommand.NormalizeSample(42));
+    }
+
+    [Fact]
+    public void InferSchema_ReportsTopLevelTypes()
+    {
+        var documents = Parse(
+            "{\"id\":\"a\",\"price\":10,\"active\":true}",
+            "{\"id\":\"b\",\"price\":20,\"active\":false}");
+
+        var fields = SchemaCommand.InferSchema(documents);
+
+        Assert.Equal(new[] { "active", "id", "price" }, fields.Select(f => f.Path));
+        Assert.Equal(new[] { "string" }, Field(fields, "id").Types);
+        Assert.Equal(new[] { "number" }, Field(fields, "price").Types);
+        Assert.Equal(new[] { "boolean" }, Field(fields, "active").Types);
+    }
+
+    [Fact]
+    public void InferSchema_TracksPresenceAcrossDocuments()
+    {
+        var documents = Parse(
+            "{\"id\":\"a\",\"optional\":1}",
+            "{\"id\":\"b\"}",
+            "{\"id\":\"c\"}");
+
+        var fields = SchemaCommand.InferSchema(documents);
+
+        Assert.Equal(3, Field(fields, "id").Presence);
+        Assert.Equal(1, Field(fields, "optional").Presence);
+    }
+
+    [Fact]
+    public void InferSchema_RecordsMultipleTypesForSameField()
+    {
+        var documents = Parse(
+            "{\"value\":1}",
+            "{\"value\":\"text\"}",
+            "{\"value\":null}");
+
+        var types = Field(SchemaCommand.InferSchema(documents), "value").Types;
+
+        Assert.Equal(new[] { "null", "number", "string" }, types);
+    }
+
+    [Fact]
+    public void InferSchema_DescribesNestedObjectsWithDotNotation()
+    {
+        var documents = Parse("{\"address\":{\"city\":\"Seattle\",\"zip\":\"98101\"}}");
+
+        var fields = SchemaCommand.InferSchema(documents);
+
+        Assert.Equal(new[] { "object" }, Field(fields, "address").Types);
+        Assert.Equal(new[] { "string" }, Field(fields, "address.city").Types);
+        Assert.Equal(new[] { "string" }, Field(fields, "address.zip").Types);
+    }
+
+    [Fact]
+    public void InferSchema_HonorsMaxDepth()
+    {
+        var documents = Parse("{\"a\":{\"b\":{\"c\":1}}}");
+
+        var fields = SchemaCommand.InferSchema(documents, maxDepth: 1);
+
+        Assert.Contains(fields, f => f.Path == "a");
+        Assert.DoesNotContain(fields, f => f.Path == "a.b");
+    }
+
+    [Fact]
+    public void InferSchema_ReportsArrayType()
+    {
+        var documents = Parse("{\"tags\":[1,2,3]}");
+
+        Assert.Equal(new[] { "array" }, Field(SchemaCommand.InferSchema(documents), "tags").Types);
+    }
+
+    [Fact]
+    public void InferSchema_IgnoresNonObjectDocuments()
+    {
+        var documents = Parse("42", "\"text\"", "{\"id\":\"a\"}");
+
+        var fields = SchemaCommand.InferSchema(documents);
+
+        Assert.Single(fields);
+        Assert.Equal("id", fields[0].Path);
+    }
+
+    private static SchemaCommand.FieldSchema Field(IReadOnlyList<SchemaCommand.FieldSchema> fields, string path)
+    {
+        return fields.Single(f => f.Path == path);
+    }
+
+    private static List<JsonElement> Parse(params string[] json)
+    {
+        var documents = new List<JsonElement>(json.Length);
+        foreach (var text in json)
+        {
+            using var document = JsonDocument.Parse(text);
+            documents.Add(document.RootElement.Clone());
+        }
+
+        return documents;
+    }
+}

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SchemaCommand.cs

@@ -0,0 +1,234 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Azure.Data.Cosmos.Shell.Commands;
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text.Json;
+using Azure.Data.Cosmos.Shell.Mcp;
+using Azure.Data.Cosmos.Shell.Parser;
+using Azure.Data.Cosmos.Shell.Util;
+using global::Azure.Data.Cosmos.Shell.Core;
+using global::Azure.Data.Cosmos.Shell.States;
+
+[CosmosCommand("schema")]
+[CosmosExample("schema", Description = "Infer the schema of the current container from a small sample")]
+[CosmosExample("schema --sample=50", Description = "Sample up to 50 documents when inferring the schema")]
+[CosmosExample("schema --database=MyDB --container=Products", Description = "Infer the schema for a specific database and container")]
+[McpAnnotation(
+    Title = "Schema",
+    ReadOnly = true,
+    Idempotent = true,
+    OpenWorld = true,
+    Description = "Returns a cheap, bounded discovery summary of a Cosmos DB container: partition key path(s), an indexing policy summary, an estimated document count, and inferred field types from a bounded sample of documents. Use this before querying to avoid re-sampling and to avoid guessing field names.")]
+internal class SchemaCommand : CosmosCommand
+{
+    internal const int DefaultSample = 20;
+    internal const int MinSample = 1;
+    internal const int MaxSample = 100;
+
+    private const int DefaultMaxDepth = 8;
+    private const string ResourceUsageHeader = "x-ms-resource-usage";
+
+    [CosmosOption("database", "db")]
+    public string? Database { get; init; }
+
+    [CosmosOption("container", "con")]
+    public string? Container { get; init; }
+
+    [CosmosOption("sample", "s")]
+    public int? Sample { get; init; }
+
+    public async override Task<CommandState> ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token)
+    {
+        if (shell.State is not ConnectedState connectedState)
+        {
+            throw new NotConnectedException("schema");
+        }
+
+        var (databaseName, containerName, container) = await ResolveContainerAsync(
+            connectedState.Client,
+            shell.State,
+            this.Database,
+            this.Container,
+            "schema",
+            token);
+
+        int sampleSize = NormalizeSample(this.Sample);
+
+        try
+        {
+            var settings = await CosmosResourceFacade.GetContainerSettingsAsync(connectedState, databaseName!, containerName!, token);
+            var sampledDocuments = await SampleDocumentsAsync(container, sampleSize, token);
+            long? documentCountEstimate = await ReadDocumentCountEstimateAsync(container, token);
+            var fields = InferSchema(sampledDocuments);
+
+            return BuildResult(databaseName!, containerName!, settings, documentCountEstimate, sampleSize, sampledDocuments.Count, fields);
+        }
+        catch (Exception e) when (e is not OperationCanceledException)
+        {
+            throw new CommandException("schema", e);
+        }
+    }
+
+    /// <summary>
+    /// Clamps the requested sample size into the supported <see cref="MinSample"/>..<see cref="MaxSample"/>
+    /// range so the discovery query stays bounded regardless of the value supplied.
+    /// </summary>
+    internal static int NormalizeSample(int? sample)
+    {
+        if (!sample.HasValue)
+        {
+            return DefaultSample;
+        }
+
+        return Math.Clamp(sample.Value, MinSample, MaxSample);
+    }
+
+    /// <summary>
+    /// Infers a per-field type summary from a bounded set of sampled documents. Fields are
+    /// reported using dot notation for nested objects up to <paramref name="maxDepth"/> levels.
+    /// Each field lists the distinct JSON types observed and the number of sampled documents in
+    /// which the field was present.
+    /// </summary>
+    internal static IReadOnlyList<FieldSchema> InferSchema(IReadOnlyList<JsonElement> documents, int maxDepth = DefaultMaxDepth)
+    {
+        var fields = new Dictionary<string, FieldAccumulator>(StringComparer.Ordinal);
+
+        foreach (var document in documents)
+        {
+            if (document.ValueKind != JsonValueKind.Object)
+            {
+                continue;
+            }
+
+            CollectFields(document, prefix: string.Empty, depth: 0, maxDepth, fields);
+        }
+
+        return fields
+            .OrderBy(pair => pair.Key, StringComparer.Ordinal)
+            .Select(pair => new FieldSchema(pair.Key, pair.Value.Types.ToArray(), pair.Value.Presence))
+            .ToList();
+    }
+
+    private static void CollectFields(JsonElement element, string prefix, int depth, int maxDepth, Dictionary<string, FieldAccumulator> fields)
+    {
+        foreach (var property in element.EnumerateObject())
+        {
+            string path = prefix.Length == 0 ? property.Name : $"{prefix}.{property.Name}";
+
+            if (!fields.TryGetValue(path, out var accumulator))
+            {
+                accumulator = new FieldAccumulator();
+                fields[path] = accumulator;
+            }
+
+            accumulator.Types.Add(DescribeValueKind(property.Value.ValueKind));
+            accumulator.Presence++;
+
+            if (property.Value.ValueKind == JsonValueKind.Object && depth + 1 < maxDepth)
+            {
+                CollectFields(property.Value, path, depth + 1, maxDepth, fields);
+            }
+        }
+    }
+
+    private static string DescribeValueKind(JsonValueKind kind) => kind switch
+    {
+        JsonValueKind.String => "string",
+        JsonValueKind.Number => "number",
+        JsonValueKind.True or JsonValueKind.False => "boolean",
+        JsonValueKind.Object => "object",
+        JsonValueKind.Array => "array",
+        JsonValueKind.Null => "null",
+        _ => "undefined",
+    };
+
+    private static async Task<List<JsonElement>> SampleDocumentsAsync(Container container, int sample, CancellationToken token)
+    {
+        var documents = new List<JsonElement>(sample);
+        using var iterator = container.GetItemQueryIterator<JsonElement>(
+            new QueryDefinition("SELECT * FROM c"),
+            requestOptions: new QueryRequestOptions { MaxItemCount = sample });
+
+        while (iterator.HasMoreResults && documents.Count < sample)
+        {
+            foreach (var element in await iterator.ReadNextAsync(token))
+            {
+                documents.Add(element.Clone());
+                if (documents.Count >= sample)
+                {
+                    break;
+                }
+            }
+        }
+
+        return documents;
+    }
+
+    private static async Task<long?> ReadDocumentCountEstimateAsync(Container container, CancellationToken token)
+    {
+        var response = await container.ReadContainerAsync(new ContainerRequestOptions { PopulateQuotaInfo = true }, token);
+        return InfoCommand.ParseResourceUsage(response.Headers[ResourceUsageHeader]).DocumentCount;
+    }
+
+    private static CommandState BuildResult(
+        string databaseName,
+        string containerName,
+        ContainerSettingsView settings,
+        long? documentCountEstimate,
+        int sampleSize,
+        int sampledDocuments,
+        IReadOnlyList<FieldSchema> fields)
+    {
+        Dictionary<string, object?>? indexingPolicy = settings.IndexingPolicy is { } indexing
+            ? new Dictionary<string, object?>
+            {
+                ["indexingMode"] = indexing.IndexingMode,
+                ["automatic"] = indexing.Automatic,
+                ["includedPaths"] = indexing.IncludedPathCount,
+                ["excludedPaths"] = indexing.ExcludedPathCount,
+                ["compositeIndexes"] = indexing.CompositeIndexCount,
+                ["spatialIndexes"] = indexing.SpatialIndexCount,
+                ["vectorIndexes"] = indexing.VectorIndexCount,
+            }
+            : null;
+
+        var output = new Dictionary<string, object?>
+        {
+            ["database"] = databaseName,
+            ["container"] = containerName,
+            ["partitionKeyPaths"] = settings.PartitionKeyPaths,
+            ["documentCountEstimate"] = documentCountEstimate,
+            ["sampleSize"] = sampleSize,
+            ["sampledDocuments"] = sampledDocuments,
+            ["indexingPolicy"] = indexingPolicy,
+            ["fields"] = fields.Select(field => new Dictionary<string, object?>
+            {
+                ["path"] = field.Path,
+                ["types"] = field.Types,
+                ["presence"] = field.Presence,
+            }).ToList(),
+        };
+
+        return new CommandState
+        {
+            Result = new ShellJson(JsonSerializer.SerializeToElement(output)),
+        };
+    }
+
+    /// <summary>
+    /// The inferred type summary for a single field discovered while sampling a container.
+    /// </summary>
+    internal sealed record FieldSchema(string Path, IReadOnlyList<string> Types, int Presence);
+
+    private sealed class FieldAccumulator
+    {
+        public SortedSet<string> Types { get; } = new(StringComparer.Ordinal);
+
+        public int Presence { get; set; }
+    }
+}