Fix #3317 and #3322 add after_key support to composite aggregation re… (#3367)

Fix #3317 and #3322 add after_key support to composite aggregation result and allow format to be set for date histogram composite source

Introduce CompositeBucketAggregate derived from  MultiBucketAggregate<CompositeBucket> to introduce AfterKey property.
Bump dotnet-xunit .NET tool to official release

Author: Mpdreamz

src/Nest/Aggregations/Aggregate.cs

@@ -3,12 +3,15 @@
 namespace Nest
 {
 	/// <summary>
-	/// Represents the result of an aggregation on the response
+	/// Aggregation response for an aggregation request
 	/// </summary>
 	[ExactContractJsonConverter(typeof(AggregateJsonConverter))]
 	public interface IAggregate
 	{
 		//TODO this public set is problematic
+		/// <summary>
+		/// Metadata for the aggregation
+		/// </summary>
 		IReadOnlyDictionary<string, object> Meta { get; set; }
 	}
 }

src/Nest/Aggregations/AggregateDictionary.cs

@@ -175,7 +175,17 @@ public TermsAggregate<TKey> Terms<TKey>(string key)
 
 		public MultiBucketAggregate<DateHistogramBucket> DateHistogram(string key) => GetMultiBucketAggregate<DateHistogramBucket>(key);
 
-		public MultiBucketAggregate<CompositeBucket> Composite(string key) => GetMultiBucketAggregate<CompositeBucket>(key);
+		public CompositeBucketAggregate Composite(string key)
+		{
+			var bucket = this.TryGet<BucketAggregate>(key);
+			if (bucket == null) return null;
+			return new CompositeBucketAggregate
+			{
+				Buckets = bucket.Items.OfType<CompositeBucket>().ToList(),
+				Meta = bucket.Meta,
+				AfterKey = new CompositeKey(bucket.AfterKey)
+			};
+		}
 
 		public MatrixStatsAggregate MatrixStats(string key) => this.TryGet<MatrixStatsAggregate>(key);
 

src/Nest/Aggregations/AggregateJsonConverter.cs

@@ -36,6 +36,7 @@ private static class Parser
 			public const string Hits = "hits";
 			public const string Location = "location";
 			public const string Fields = "fields";
+			public const string AfterKey = "after_key";
 
 			public const string Key = "key";
 			public const string From = "from";
@@ -126,6 +127,16 @@ private IAggregate ReadAggregate(JsonReader reader, JsonSerializer serializer)
 				case Parser.Value:
 					aggregate = GetValueAggregate(reader, serializer);
 					break;
+				case Parser.AfterKey:
+					reader.Read();
+					var afterKeys = serializer.Deserialize<Dictionary<string, object>>(reader);
+					reader.Read();
+					var bucketAggregate = reader.Value.ToString() == Parser.Buckets
+						? this.GetMultiBucketAggregate(reader, serializer) as BucketAggregate ?? new BucketAggregate()
+						: new BucketAggregate();
+					bucketAggregate.AfterKey = afterKeys;
+					aggregate = bucketAggregate;
+					break;
 				case Parser.Buckets:
 				case Parser.DocCountErrorUpperBound:
 					aggregate = GetMultiBucketAggregate(reader, serializer);

src/Nest/Aggregations/Bucket/BucketAggregate.cs

@@ -23,17 +23,42 @@ public SingleBucketAggregate(IReadOnlyDictionary<string, IAggregate> aggregation
 		[Obsolete("Use methods on this instance to access sub aggregations. Will be removed in NEST 7.x")]
 		public AggregateDictionary Aggregations { get; protected internal set; }
 
+		/// <summary>
+		/// Count of documents in the bucket
+		/// </summary>
 		public long DocCount { get; internal set; }
 	}
 
+	/// <summary>
+	/// Aggregation response for a bucket aggregation
+	/// </summary>
+	/// <typeparam name="TBucket"></typeparam>
 	public class MultiBucketAggregate<TBucket> : IAggregate
 		where TBucket : IBucket
 	{
+		/// <inheritdoc />
 		public IReadOnlyDictionary<string, object> Meta { get; set; }
 
+		/// <summary>
+		/// The buckets into which results are grouped
+		/// </summary>
 		public IReadOnlyCollection<TBucket> Buckets { get; set; } = EmptyReadOnly<TBucket>.Collection;
 	}
 
+	/// <summary>
+	/// Aggregation response of <see cref="CompositeAggregation"/>
+	/// </summary>
+	public class CompositeBucketAggregate : MultiBucketAggregate<CompositeBucket>
+	{
+		/// <summary>
+		/// The composite key of the last bucket returned
+		/// in the response before any filtering by pipeline aggregations.
+		/// If all buckets are filtered/removed by pipeline aggregations,
+		/// <see cref="AfterKey"/> will contain the composite key of the last bucket before filtering.
+		/// </summary>
+		/// <remarks> Valid for Elasticsearch 6.3.0+ </remarks>
+		public CompositeKey AfterKey { get; set; }
+	}
 
 	// Intermediate object used for deserialization
 	public class BucketAggregate : IAggregate
@@ -44,5 +69,6 @@ public class BucketAggregate : IAggregate
 		public IReadOnlyDictionary<string, object> Meta { get; set; } = EmptyReadOnly<string, object>.Dictionary;
 		public long DocCount { get; set; }
 		public long BgCount { get; set; }
+		public IReadOnlyDictionary<string, object> AfterKey { get; set; } = EmptyReadOnly<string, object>.Dictionary;
 	}
 }

src/Nest/Aggregations/Bucket/Composite/DateHistogramCompositeAggregationSource.cs

@@ -22,6 +22,13 @@ public interface IDateHistogramCompositeAggregationSource : ICompositeAggregatio
 		/// </summary>
 		[JsonProperty("time_zone")]
 		string Timezone { get; set; }
+
+		/// <summary>
+		/// Return a formatted date string as the key instead an epoch long
+		/// </summary>
+		/// <remarks> Valid for Elasticsearch 6.3.0+ </remarks>
+		[JsonProperty("format")]
+		string Format { get; set; }
 	}
 
 	/// <inheritdoc cref="IDateHistogramCompositeAggregationSource"/>
@@ -35,6 +42,9 @@ public DateHistogramCompositeAggregationSource(string name) : base(name) {}
 		/// <inheritdoc />
 		public string Timezone { get; set; }
 
+		/// <inheritdoc />
+		public string Format { get; set; }
+
 		/// <inheritdoc />
 		protected override string SourceType => "date_histogram";
 	}
@@ -46,6 +56,7 @@ public class DateHistogramCompositeAggregationSourceDescriptor<T>
 	{
 		Union<DateInterval?,Time> IDateHistogramCompositeAggregationSource.Interval { get; set; }
 		string IDateHistogramCompositeAggregationSource.Timezone { get; set; }
+		string IDateHistogramCompositeAggregationSource.Format { get; set; }
 
 		public DateHistogramCompositeAggregationSourceDescriptor(string name) : base(name, "date_histogram") {}
 
@@ -58,7 +69,9 @@ public DateHistogramCompositeAggregationSourceDescriptor<T> Interval(Time interv
 			Assign(a => a.Interval = interval);
 
 		/// <inheritdoc cref="IDateHistogramCompositeAggregationSource.Timezone"/>
-		public DateHistogramCompositeAggregationSourceDescriptor<T> Timezone(string timezone) =>
-			Assign(a => a.Timezone = timezone);
+		public DateHistogramCompositeAggregationSourceDescriptor<T> Timezone(string timezone) => Assign(a => a.Timezone = timezone);
+
+		/// <inheritdoc cref="IDateHistogramCompositeAggregationSource.Timezone"/>
+		public DateHistogramCompositeAggregationSourceDescriptor<T> Format(string format) => Assign(a => a.Format = format);
 	}
 }

src/Tests/Tests/Aggregations/Bucket/Composite/CompositeAggregationUsageTests.cs

@@ -1,10 +1,12 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using Elastic.Xunit.XunitPlumbing;
 using FluentAssertions;
 using Nest;
 using Newtonsoft.Json;
+using Tests.Configuration;
 using Tests.Core.Extensions;
 using Tests.Core.ManagedElasticsearch.Clusters;
 using Tests.Domain;
@@ -163,6 +165,12 @@ protected override void ExpectResponse(ISearchResponse<Project> response)
 			var composite = response.Aggregations.Composite("my_buckets");
 			composite.Should().NotBeNull();
 			composite.Buckets.Should().NotBeNullOrEmpty();
+			composite.AfterKey.Should().NotBeNull();
+			if (TestConfiguration.Instance.InRange(">=6.3.0"))
+			{
+				composite.AfterKey.Should().HaveCount(3)
+					.And.ContainKeys("branches", "started", "branch_count");
+			}
 			foreach (var item in composite.Buckets)
 			{
 				var key = item.Key;
@@ -187,4 +195,118 @@ protected override void ExpectResponse(ISearchResponse<Project> response)
 			}
 		}
 	}
+
+
+	//hide
+	[SkipVersion("<6.3.0", "Date histogram source only supports format starting from Elasticsearch 6.3.0+")]
+	public class DateFormatCompositeAggregationUsageTests : ProjectsOnlyAggregationUsageTestBase
+	{
+		public DateFormatCompositeAggregationUsageTests(ReadOnlyCluster i, EndpointUsage usage) : base(i, usage) { }
+
+		protected override object AggregationJson => new
+		{
+			my_buckets = new
+			{
+				composite = new
+				{
+					sources = new object[]
+					{
+						new
+						{
+							started = new
+							{
+								date_histogram = new
+								{
+									field = "startedOn",
+									interval = "month",
+									format = "yyyy-MM-dd"
+								}
+							}
+						},
+					}
+				},
+				aggs = new
+				{
+					project_tags = new
+					{
+						nested = new
+						{
+							path = "tags"
+						},
+						aggs = new
+						{
+							tags = new
+							{
+								terms = new {field = "tags.name"}
+							}
+						}
+					}
+				}
+			}
+		};
+
+		protected override Func<AggregationContainerDescriptor<Project>, IAggregationContainer> FluentAggs => a => a
+			.Composite("my_buckets", date => date
+				.Sources(s => s
+					.DateHistogram("started", d => d
+						.Field(f => f.StartedOn)
+						.Interval(DateInterval.Month)
+						.Format("yyyy-MM-dd")
+					)
+				)
+				.Aggregations(childAggs => childAggs
+					.Nested("project_tags", n => n
+						.Path(p => p.Tags)
+						.Aggregations(nestedAggs => nestedAggs
+							.Terms("tags", avg => avg.Field(p => p.Tags.First().Name))
+						)
+					)
+				)
+			);
+
+		protected override AggregationDictionary InitializerAggs =>
+			new CompositeAggregation("my_buckets")
+			{
+				Sources = new List<ICompositeAggregationSource>
+				{
+					new DateHistogramCompositeAggregationSource("started")
+					{
+						Field = Infer.Field<Project>(f => f.StartedOn),
+						Interval = DateInterval.Month,
+						Format = "yyyy-MM-dd"
+					},
+				},
+				Aggregations = new NestedAggregation("project_tags")
+				{
+					Path = Field<Project>(p => p.Tags),
+					Aggregations = new TermsAggregation("tags")
+					{
+						Field = Field<Project>(p => p.Tags.First().Name)
+					}
+				}
+			};
+
+		/**==== Handling Responses
+		 * Each Composite aggregation bucket key is an `CompositeKey`, a specialized
+		 * `IReadOnlyDictionary<string, object>` type with methods to convert values to supported types
+		 */
+		protected override void ExpectResponse(ISearchResponse<Project> response)
+		{
+			response.ShouldBeValid();
+
+			var composite = response.Aggregations.Composite("my_buckets");
+			composite.Should().NotBeNull();
+			composite.Buckets.Should().NotBeNullOrEmpty();
+			composite.AfterKey.Should().NotBeNull();
+			composite.AfterKey.Should().HaveCount(1).And.ContainKeys("started");
+			foreach (var item in composite.Buckets)
+			{
+				var key = item.Key;
+				key.Should().NotBeNull();
+
+				key.TryGetValue("started", out string startedString).Should().BeTrue();
+				startedString.Should().NotBeNullOrWhiteSpace();
+			}
+		}
+	}
 }

src/Tests/Tests/Tests.csproj

@@ -7,7 +7,7 @@
     <NoWarn>$(NoWarn);xUnit1013</NoWarn>
   </PropertyGroup>
   <ItemGroup>
-    <DotNetCliToolReference Include="dotnet-xunit" Version="2.3.0-beta1-build3642" />
+    <DotNetCliToolReference Include="dotnet-xunit" Version="2.3.1" />
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="..\Tests.Core\Tests.Core.csproj" />