Add support for Query API (#810)

halspang · rynowak · web-flow · commit d43de2b04381 · 2022-01-12T17:50:06.000-08:00
This commit adds methods to the DaprClient to call the new Query API. Currently, a raw query string is provided to the API. The API is still in the alpha stage which may cause shifts in how the API works. This should not effect how the SDK exposes the call, but may change the underlying semantics. #777 Signed-off-by: Hal Spang <halspang@microsoft.com> Co-authored-by: Ryan Nowak <nowakra@gmail.com>
diff --git a/src/Dapr.Client/DaprClient.cs b/src/Dapr.Client/DaprClient.cs
@@ -757,6 +757,22 @@ public abstract Task<bool> TryDeleteStateAsync(
             IReadOnlyDictionary<string, string> metadata = default,
             CancellationToken cancellationToken = default);
 
+        /// <summary>
+        /// Queries the specified statestore with the given query. The query is a JSON representation of the query as described by the Dapr QueryState API.
+        /// Note that the underlying statestore must support queries.
+        /// </summary>
+        /// <param name="storeName">The name of the statestore.</param>
+        /// <param name="jsonQuery">A JSON formatted query string.</param>
+        /// <param name="metadata">Metadata to send to the statestore.</param>
+        /// <param name="cancellationToken">A <see cref="CancellationToken"/> that can be used to cancel the operation.</param>
+        /// <typeparam name="TValue">The data type of the value to read.</typeparam>
+        /// <returns>A <see cref="StateQueryResponse{TValue}"/> that may be paginated, use <see cref="StateQueryResponse{TValue}.Token"/> to continue the query.</returns>
+        public abstract Task<StateQueryResponse<TValue>> QueryStateAsync<TValue>(
+            string storeName,
+            string jsonQuery,
+            IReadOnlyDictionary<string, string> metadata = default,
+            CancellationToken cancellationToken = default);
+
         /// <summary>
         /// Gets the secret value from the secret store.
         /// </summary>
diff --git a/src/Dapr.Client/DaprClientGrpc.cs b/src/Dapr.Client/DaprClientGrpc.cs
@@ -1015,6 +1015,65 @@ private async Task<bool> MakeDeleteStateCallAsync(
                 throw new DaprException("State operation failed: the Dapr endpoint indicated a failure. See InnerException for details.", ex);
             }
         }
+
+        /// <inheritdoc/>
+        public async override Task<StateQueryResponse<TValue>> QueryStateAsync<TValue>(
+            string storeName,
+            string jsonQuery,
+            IReadOnlyDictionary<string, string> metadata = default,
+            CancellationToken cancellationToken = default)
+        {
+            var queryRequest = new Autogenerated.QueryStateRequest()
+            {
+                StoreName = storeName,
+                Query = jsonQuery
+            };
+
+            if (metadata != null)
+            {
+                foreach (var kvp in metadata)
+                {
+                    queryRequest.Metadata.Add(kvp.Key, kvp.Value);
+                }
+            }
+
+            var options = CreateCallOptions(headers: null, cancellationToken);
+
+            try
+            {
+                var items = new List<StateQueryItem<TValue>>();
+                var failedKeys = new List<string>();
+                var queryResponse = await client.QueryStateAlpha1Async(queryRequest, options);
+                foreach (var item in queryResponse.Results)
+                {
+                    if (!string.IsNullOrEmpty(item.Error))
+                    {
+                        // When we encounter an error, we record the key and prepare to throw an exception at the end of the results.
+                        failedKeys.Add(item.Key);
+                        continue;
+                    }
+
+                    items.Add(new StateQueryItem<TValue>(item.Key, TypeConverters.FromJsonByteString<TValue>(item.Data, this.JsonSerializerOptions), item.Etag, item.Error));
+                }
+
+                var results = new StateQueryResponse<TValue>(items, queryResponse.Token, queryResponse.Metadata);
+                if (failedKeys.Count > 0)
+                {
+                    // We encountered some bad keys so we throw instead of returning to alert the user.
+                    throw new StateQueryException<TValue>($"Encountered an error while processing state query results.", results, failedKeys);
+                }
+
+                return results;
+            }
+            catch (RpcException ex)
+            {
+                throw new DaprException("Query state operation failed: the Dapr endpointed indicated a failure. See InnerException for details.", ex);
+            }
+            catch (JsonException ex)
+            {
+                throw new DaprException("State operation failed: the state payload could not be deserialized. See InnerException for details.", ex);
+            }
+        }
         #endregion
 
         #region Secret Apis
diff --git a/src/Dapr.Client/StateQueryException.cs b/src/Dapr.Client/StateQueryException.cs
@@ -0,0 +1,34 @@
+﻿using System.Collections.Generic;
+
+namespace Dapr.Client
+{
+    /// <summary>
+    /// Exception that is thrown when an erorr is encountered during a call to the Query API.
+    /// This exception contains the partial results (if any) from that exception.
+    /// </summary>
+    public class StateQueryException<TValue> : DaprApiException
+    {
+        /// <summary>
+        /// Constructor.
+        /// </summary>
+        /// <param name="message">The description of the exception from the source.</param>
+        /// <param name="response">The response containing successful items, if any, in a response with errors.</param>
+        /// <param name="failedKeys">The key(s) that encountered an error during the query.</param>
+        public StateQueryException(string message, StateQueryResponse<TValue> response, IReadOnlyList<string> failedKeys)
+            : base(message)
+        {
+            Response = response;
+            FailedKeys = failedKeys;
+        }
+
+        /// <summary>
+        /// The response containing successful items, if any, in a response with errors.
+        /// </summary>
+        public StateQueryResponse<TValue> Response { get; }
+
+        /// <summary>
+        /// The key(s) that encountered an error during the query.
+        /// </summary>
+        public IReadOnlyList<string> FailedKeys { get; }
+    }
+}
diff --git a/src/Dapr.Client/StateQueryResponse.cs b/src/Dapr.Client/StateQueryResponse.cs
@@ -0,0 +1,92 @@
+﻿/*
+Copyright 2021 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+using System.Collections.Generic;
+
+namespace Dapr.Client
+{
+    /// <summary>
+    /// Represents the response from a state query.
+    /// </summary>
+    public class StateQueryResponse<TValue>
+    {
+        /// <summary>
+        /// Constructor.
+        /// </summary>
+        /// <param name="results">The results of the query.</param>
+        /// <param name="token">The pagination token to continue the query.</param>
+        /// <param name="metadata">The metadata to be passed back to the caller.</param>
+        public StateQueryResponse(IReadOnlyList<StateQueryItem<TValue>> results, string token, IReadOnlyDictionary<string, string> metadata)
+        {
+            Results = new List<StateQueryItem<TValue>>(results);
+            Token = token;
+            Metadata = metadata;
+        }
+
+        /// <summary>
+        /// The results of the query.
+        /// </summary>
+        public IReadOnlyList<StateQueryItem<TValue>> Results { get; }
+
+        /// <summary>
+        /// The pagination token to continue the query.
+        /// </summary>
+        public string Token { get; }
+
+        /// <summary>
+        /// The metadata to be passed back to the caller.
+        /// </summary>
+        public IReadOnlyDictionary<string, string> Metadata { get; }
+    }
+
+    /// <summary>
+    /// Represents an individual item from the results of a state query.
+    /// </summary>
+    public class StateQueryItem<TValue>
+    {
+        /// <summary>
+        /// Constructor.
+        /// </summary>
+        /// <param name="key">The key of the returned item.</param>
+        /// <param name="data">The value of the returned item.</param>
+        /// <param name="etag">The ETag of the returned item.</param>
+        /// <param name="error">The error, if one occurred, of the returned item.</param>
+        public StateQueryItem(string key, TValue data, string etag, string error)
+        {
+            Key = key;
+            Data = data;
+            ETag = etag;
+            Error = error;
+        }
+
+        /// <summary>
+        /// The key from the matched query.
+        /// </summary>
+        public string Key { get; }
+
+        /// <summary>
+        /// The data of the the key from the matched query.
+        /// </summary>
+        public TValue Data { get; }
+
+        /// <summary>
+        /// The ETag for the key from the matched query.
+        /// </summary>
+        public string ETag { get; }
+
+        /// <summary>
+        /// The error from the query, if one occurred.
+        /// </summary>
+        public string Error { get; }
+    }
+}
diff --git a/test/Dapr.Client.Test/StateApiTest.cs b/test/Dapr.Client.Test/StateApiTest.cs
@@ -29,6 +29,7 @@ namespace Dapr.Client.Test
     using Xunit;
     using System.Threading;
     using System.Net.Http;
+    using System.Text;
 
     public class StateApiTest
     {
@@ -999,6 +1000,81 @@ public async Task DeleteBulkStateAsync_ValidateRequest()
             envelope.States[0].Metadata.Should().ContainKey("partitionKey");
         }
 
+        [Fact]
+        public async Task QueryStateAsync_ValidateResult()
+        {
+            await using var client = TestClient.CreateForDaprClient();
+
+            var queryJson = "{'query':{'filter':{ 'EQ': {'value':'test'}}}}";
+            var request = await client.CaptureGrpcRequestAsync(async daprClient =>
+            {
+                return await daprClient.QueryStateAsync<Widget>("testStore", queryJson, new Dictionary<string, string>());
+            });
+
+            // Validate request.
+            var envelope = await request.GetRequestEnvelopeAsync<Autogenerated.QueryStateRequest>();
+            envelope.StoreName.Should().Be("testStore");
+            envelope.Query.Should().Be(queryJson);
+            envelope.Metadata.Should().BeEmpty();
+
+            // Validate response.
+            var testData = new Widget() { Color = "Green", Size = "Small" };
+            var wireResponse = new Autogenerated.QueryStateResponse();
+            wireResponse.Results.Add(MakeQueryStateItem("test", testData, "an etag"));
+
+            var response = await request.CompleteWithMessageAsync(wireResponse);
+            response.Results.Count.Should().Be(1);
+            response.Results[0].Key.Should().Be("test");
+            response.Results[0].Data.Should().Be(testData);
+            response.Results[0].ETag.Should().Be("an etag");
+            response.Results[0].Error.Should().BeNullOrEmpty();
+        }
+
+        [Fact]
+        public async Task QueryStateAsync_EncountersError_ValidatePartialResult()
+        {
+            await using var client = TestClient.CreateForDaprClient();
+
+            var queryJson = "{'query':{'filter':{ 'EQ': {'value':'test'}}}}";
+            var request = await client.CaptureGrpcRequestAsync(async daprClient =>
+            {
+                return await daprClient.QueryStateAsync<Widget>("testStore", queryJson, new Dictionary<string, string>());
+            });
+
+            // Validate request.
+            var envelope = await request.GetRequestEnvelopeAsync<Autogenerated.QueryStateRequest>();
+            envelope.StoreName.Should().Be("testStore");
+            envelope.Query.Should().Be(queryJson);
+            envelope.Metadata.Should().BeEmpty();
+
+            // Validate response, we expect to only get the first object as the 2nd will present an error.
+            var testData1 = new Widget() { Color = "Green", Size = "Small" };
+            var testData2 = new Widget() { Color = "Green", Size = "Medium" };
+            var testData3 = new Widget() { Color = "Green", Size = "Large" };
+            var wireResponse = new Autogenerated.QueryStateResponse();
+
+            wireResponse.Results.Add(MakeQueryStateItem("test1", testData1));
+            wireResponse.Results.Add(MakeQueryStateItem("test2", testData2, string.Empty, "An error!"));
+            wireResponse.Results.Add(MakeQueryStateItem("test3", testData3));
+
+            var ex = await Assert.ThrowsAsync<StateQueryException<Widget>>(() => request.CompleteWithMessageAsync(wireResponse));
+            ex.Message.Should().Be("Encountered an error while processing state query results.");
+            var response = ex.Response;
+            response.Results.Count.Should().Be(2);
+            response.Results[0].Key.Should().Be("test1");
+            response.Results[0].Data.Should().Be(testData1);
+            response.Results[0].ETag.Should().BeNullOrEmpty();
+            response.Results[0].Error.Should().BeNullOrEmpty();
+            response.Results[1].Key.Should().Be("test3");
+            response.Results[1].Data.Should().Be(testData3);
+            response.Results[1].ETag.Should().BeNullOrEmpty();
+            response.Results[1].Error.Should().BeNullOrEmpty();
+
+            var failedKeys = ex.FailedKeys;
+            failedKeys.Count.Should().Be(1);
+            failedKeys[0].Should().Be("test2");
+        }
+
         private Autogenerated.GetStateResponse MakeGetStateResponse<T>(T state, string etag = null)
         {
             var data = TypeConverters.ToJsonByteString(state, new JsonSerializerOptions(JsonSerializerDefaults.Web));
@@ -1033,11 +1109,33 @@ private Autogenerated.GetBulkStateResponse MakeGetBulkStateResponse<T>(string ke
             return response;
         }
 
+        private Autogenerated.QueryStateItem MakeQueryStateItem<T>(string key, T data, string etag = default, string error = default)
+        {
+            var wireItem = new Autogenerated.QueryStateItem();
+            wireItem.Key = key;
+            wireItem.Data = ByteString.CopyFromUtf8(JsonSerializer.Serialize(data));
+            wireItem.Etag = etag ?? string.Empty;
+            wireItem.Error = error ?? string.Empty;
+            return wireItem;
+        }
+
         private class Widget
         {
             public string Size { get; set; }
 
             public string Color { get; set; }
+
+            public override bool Equals(object obj)
+            {
+                return obj is Widget widget &&
+                       Size == widget.Size &&
+                       Color == widget.Color;
+            }
+
+            public override int GetHashCode()
+            {
+                return HashCode.Combine(Size, Color);
+            }
         }
     }
 }
