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

Author: Mpdreamz

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 = 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

@@ -34,6 +34,20 @@ public class MultiBucketAggregate<TBucket> : IAggregate
 		public IReadOnlyCollection<TBucket> Buckets { get; set; } = EmptyReadOnly<TBucket>.Collection;
 	}
 
+	public class CompositeBucketAggregate : IAggregate
+	{
+		public IReadOnlyDictionary<string, object> Meta { get; set; }
+
+		public IReadOnlyCollection<CompositeBucket> Buckets { get; set; } = EmptyReadOnly<CompositeBucket>.Collection;
+
+		/// <summary>
+		/// The after_key is equals to the last bucket returned in the response before any filtering that could be done by Pipeline aggregations.
+		/// If all buckets are filtered/removed by a pipeline aggregation, the after_key will contain the last bucket before filtering.
+		/// </summary>
+		/// <remarks> Valid for Elasticsearch 6.3.0+ </remarks>
+		public IReadOnlyDictionary<string, object> AfterKey { get; set; } = EmptyReadOnly<string, object>.Dictionary;
+	}
+
 
 	// Intermediate object used for deserialization
 	public class BucketAggregate : IAggregate
@@ -44,5 +58,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/Framework/EndpointTests/ApiIntegrationTestBase.cs

@@ -32,13 +32,13 @@ protected ApiIntegrationTestBase(TCluster cluster, EndpointUsage usage) : base(c
 		public override IElasticClient Client => this.Cluster.Client;
 		protected override TInitializer Initializer => Activator.CreateInstance<TInitializer>();
 
-		[I] public async Task ReturnsExpectedStatusCode() =>
+		[I] public virtual async Task ReturnsExpectedStatusCode() =>
 			await this.AssertOnAllResponses(r => r.ApiCall.HttpStatusCode.Should().Be(this.ExpectStatusCode));
 
-		[I] public async Task ReturnsExpectedIsValid() =>
+		[I] public virtual async Task ReturnsExpectedIsValid() =>
 			await this.AssertOnAllResponses(r => r.ShouldHaveExpectedIsValid(this.ExpectIsValid));
 
-		[I] public async Task ReturnsExpectedResponse() => await this.AssertOnAllResponses(ExpectResponse);
+		[I] public virtual async Task ReturnsExpectedResponse() => await this.AssertOnAllResponses(ExpectResponse);
 
 		protected override Task AssertOnAllResponses(Action<TResponse> assert)
 		{