ls: push max limit down with SELECT TOP when no client-side filter

CosmosDBShell.Tests/CommandTests/ListCommandTests.cs

@@ -12,8 +12,85 @@ namespace CosmosShell.Tests.CommandTests;
 using Azure.Data.Cosmos.Shell.Util;
 using Microsoft.Azure.Cosmos;
 
+/// <summary>
+/// Unit tests for <see cref="ListCommand"/>. Covers <see cref="ListCommand.BuildItemQueryText"/>
+/// (which pushes the per-request limit down to the server with <c>SELECT TOP n</c>
+/// when no client-side filter is in play), the hierarchical-partition-key
+/// helpers on <see cref="CosmosCommand"/>, and the <c>ls</c>-specific
+/// <see cref="ListCommand.ReadQueryResponseAsync"/> error mapping.
+/// </summary>
 public class ListCommandTests
 {
+    [Fact]
+    public void BuildItemQueryText_NoLimit_NoFilter_UsesUnbounded()
+    {
+        Assert.Equal("SELECT * FROM c", ListCommand.BuildItemQueryText(null, null));
+    }
+
+    [Fact]
+    public void BuildItemQueryText_WithLimit_NoFilter_UsesTop()
+    {
+        Assert.Equal("SELECT TOP 25 * FROM c", ListCommand.BuildItemQueryText(25, null));
+    }
+
+    [Fact]
+    public void BuildItemQueryText_WithLimit_WildcardFilter_UsesTop()
+    {
+        // '*' is treated by ListCommand as "no filtering", so it is safe to
+        // push the cap to the server.
+        Assert.Equal("SELECT TOP 10 * FROM c", ListCommand.BuildItemQueryText(10, "*"));
+    }
+
+    [Fact]
+    public void BuildItemQueryText_WithLimit_SubstringFilter_StaysUnbounded()
+    {
+        // A substring filter is applied in the shell against the partition or
+        // custom key, so capping server-side rows would silently drop matching
+        // items. Keep paging client-side.
+        Assert.Equal("SELECT * FROM c", ListCommand.BuildItemQueryText(10, "active"));
+    }
+
+    [Fact]
+    public void BuildItemQueryText_NoLimit_SubstringFilter_StaysUnbounded()
+    {
+        Assert.Equal("SELECT * FROM c", ListCommand.BuildItemQueryText(null, "active"));
+    }
+
+    [Fact]
+    public void BuildItemQueryText_EmptyFilter_TreatedAsNoFilter()
+    {
+        Assert.Equal("SELECT TOP 5 * FROM c", ListCommand.BuildItemQueryText(5, string.Empty));
+    }
+
+    [Theory]
+    [InlineData(null, false)]
+    [InlineData("", false)]
+    [InlineData("*", false)]
+    [InlineData("active", true)]
+    public void HasClientSideFilter_ClassifiesFilter(string? filter, bool expected)
+    {
+        Assert.Equal(expected, ListCommand.HasClientSideFilter(filter));
+    }
+
+    [Fact]
+    public void ShouldReportLimitReached_ServerSideTopReportsWhenCapHitWithoutContinuation()
+    {
+        Assert.True(ListCommand.ShouldReportLimitReached(currentCount: 10, effectiveMaxItemCount: 10, usesServerSideTop: true, iteratorHasMoreResults: false));
+    }
+
+    [Fact]
+    public void ShouldReportLimitReached_ClientSideQueryRequiresContinuation()
+    {
+        Assert.False(ListCommand.ShouldReportLimitReached(currentCount: 10, effectiveMaxItemCount: 10, usesServerSideTop: false, iteratorHasMoreResults: false));
+        Assert.True(ListCommand.ShouldReportLimitReached(currentCount: 10, effectiveMaxItemCount: 10, usesServerSideTop: false, iteratorHasMoreResults: true));
+    }
+
+    [Fact]
+    public void ShouldReportLimitReached_BelowCapDoesNotReport()
+    {
+        Assert.False(ListCommand.ShouldReportLimitReached(currentCount: 9, effectiveMaxItemCount: 10, usesServerSideTop: true, iteratorHasMoreResults: true));
+    }
+
     [Fact]
     public void GetPartitionKeyPropertyNames_ReturnsAllHierarchicalPaths()
     {

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

@@ -44,7 +44,7 @@ internal class ListCommand : CosmosCommand, IStateVisitor<CommandState, ShellInt
 
     public async override Task<CommandState> ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token)
     {
-        this.matcher = string.IsNullOrEmpty(this.Filter) ? null : new PatternMatcher(this.Filter);
+        this.matcher = HasClientSideFilter(this.Filter) ? new PatternMatcher(this.Filter!) : null;
         return await shell.State.AcceptAsync(this, shell, token) ?? new CommandState();
     }
 
@@ -165,7 +165,9 @@ private async Task<CommandState> ListContainerItemsAsync(CosmosClient client, st
         var partitionKeyPropertyNames = GetPartitionKeyPropertyNames(containerResponse.Resource.PartitionKeyPaths);
         var matchKeyPropertyNames = string.IsNullOrEmpty(this.Key) ? partitionKeyPropertyNames : [this.Key];
 
-        using var feedIterator = container.GetItemQueryStreamIterator("SELECT * FROM c", requestOptions: opt);
+        var queryText = BuildItemQueryText(effectiveMaxItemCount, this.Filter);
+        var usesServerSideTop = effectiveMaxItemCount.HasValue && !HasClientSideFilter(this.Filter);
+        using var feedIterator = container.GetItemQueryStreamIterator(queryText, requestOptions: opt);
         var returnState = new CommandState();
         returnState.SetFormat(this.OutputFormat ?? Environment.GetEnvironmentVariable("COSMOSDB_SHELL_FORMAT"));
         var list = new List<JsonElement>();
@@ -178,7 +180,7 @@ private async Task<CommandState> ListContainerItemsAsync(CosmosClient client, st
             foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray())
             {
                 // Check if pattern matches
-                bool shouldList = this.matcher == null || this.Filter == "*"; // No filter or wildcard = list all
+                bool shouldList = this.matcher == null;
 
                 shouldList = shouldList || MatchesAnyPath(element, matchKeyPropertyNames, this.matcher!);
 
@@ -189,7 +191,7 @@ private async Task<CommandState> ListContainerItemsAsync(CosmosClient client, st
 
                 if (ResultLimit.IsLimitReached(list.Count, effectiveMaxItemCount))
                 {
-                    limitReached = feedIterator.HasMoreResults;
+                    limitReached = ShouldReportLimitReached(list.Count, effectiveMaxItemCount, usesServerSideTop, feedIterator.HasMoreResults);
                     break;
                 }
             }
@@ -254,4 +256,34 @@ private bool IsMatch(string item)
     {
         return this.matcher == null || this.matcher.Match(item);
     }
+
+    /// <summary>
+    /// Builds the SQL text for listing items in a container. When the caller
+    /// supplies a finite limit and there is no client-side filter that could
+    /// discard rows, switches to <c>SELECT TOP n * FROM c</c> so the server
+    /// stops scanning once the requested number of rows has been produced.
+    /// With a filter the server cannot honor the cap (the substring match is
+    /// applied in the shell against the partition or custom key), so we fall
+    /// back to <c>SELECT * FROM c</c> and rely on the existing client-side
+    /// break to stop paging.
+    /// </summary>
+    internal static string BuildItemQueryText(int? effectiveMaxItemCount, string? filter)
+    {
+        if (effectiveMaxItemCount.HasValue && !HasClientSideFilter(filter))
+        {
+            return $"SELECT TOP {effectiveMaxItemCount.Value} * FROM c";
+        }
+
+        return "SELECT * FROM c";
+    }
+
+    internal static bool HasClientSideFilter(string? filter)
+    {
+        return !string.IsNullOrEmpty(filter) && filter != "*";
+    }
+
+    internal static bool ShouldReportLimitReached(int currentCount, int? effectiveMaxItemCount, bool usesServerSideTop, bool iteratorHasMoreResults)
+    {
+        return ResultLimit.IsLimitReached(currentCount, effectiveMaxItemCount) && (usesServerSideTop || iteratorHasMoreResults);
+    }
 }