Fix time shift metric lookup (#119774) (#120028)

* fix time split metric lookup

* fix in a stable way

* add documentation

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

Co-authored-by: Joe Reuter <johannes.reuter@elastic.co>

src/plugins/data/common/search/aggs/agg_config.ts

@@ -402,6 +402,10 @@ export class AggConfig {
     return this.type.getValue(this, bucket);
   }
 
+  getResponseId() {
+    return this.type.getResponseId(this);
+  }
+
   getKey(bucket: any, key?: string) {
     if (this.type.getKey) {
       return this.type.getKey(bucket, key, this);

src/plugins/data/common/search/aggs/agg_type.ts

@@ -57,6 +57,7 @@ export interface AggTypeConfig<
   getValue?: (agg: TAggConfig, bucket: any) => any;
   getKey?: (bucket: any, key: any, agg: TAggConfig) => any;
   getValueBucketPath?: (agg: TAggConfig) => string;
+  getResponseId?: (agg: TAggConfig) => string;
 }
 
 // TODO need to make a more explicit interface for this
@@ -220,6 +221,25 @@ export class AggType<
     return false;
   }
 
+  /**
+   * Returns the key of the object containing the results of the agg in the Elasticsearch response object.
+   * In most cases this returns the `agg.id` property, but in some cases the response object is structured differently.
+   * In the following example of a terms agg, `getResponseId` returns "myAgg":
+   * ```
+   * {
+   *    "aggregations": {
+   *      "myAgg": {
+   *        "doc_count_error_upper_bound": 0,
+   *        "sum_other_doc_count": 0,
+   *        "buckets": [
+   * ...
+   * ```
+   *
+   * @param  {agg} agg - the agg to return the id in the ES reponse object for
+   * @return {string}
+   */
+  getResponseId: (agg: TAggConfig) => string;
+
   /**
    * Generic AggType Constructor
    *
@@ -293,5 +313,7 @@ export class AggType<
       });
 
     this.getValue = config.getValue || ((agg: TAggConfig, bucket: any) => {});
+
+    this.getResponseId = config.getResponseId || ((agg: TAggConfig) => agg.id);
   }
 }

src/plugins/data/common/search/aggs/metrics/filtered_metric.test.ts

@@ -69,4 +69,10 @@ describe('filtered metric agg type', () => {
       })
     ).toEqual(10);
   });
+
+  it('provides the id of the inner filter bucket to look up the agg config in the response object', () => {
+    const agg = aggConfigs.getResponseAggs()[0];
+
+    expect(agg.getResponseId()).toEqual('filtered_metric-bucket');
+  });
 });

src/plugins/data/common/search/aggs/metrics/filtered_metric.ts

@@ -52,5 +52,8 @@ export const getFilteredMetricAgg = () => {
       }
       return `${customBucket.getValueBucketPath()}>${customMetric.getValueBucketPath()}`;
     },
+    getResponseId(agg) {
+      return agg.params.customBucket.id;
+    },
   });
 };

src/plugins/data/common/search/aggs/utils/time_splits.ts

@@ -186,7 +186,7 @@ export function mergeTimeShifts(
         return;
       } else {
         // a sub-agg
-        const agg = requestAggs.find((requestAgg) => key.indexOf(requestAgg.id) === 0);
+        const agg = requestAggs.find((requestAgg) => key === requestAgg.getResponseId());
         if (agg && agg.type.type === AggGroupNames.Metrics) {
           const timeShift = agg.getTimeShift();
           if (

test/interpreter_functional/test_suites/run_pipeline/esaggs_timeshift.ts

@@ -22,6 +22,15 @@ function getCell(esaggsResult: any, row: number, column: number): unknown | unde
 function checkShift(rows: Datatable['rows'], columns: Datatable['columns'], metricIndex = 1) {
   rows.shift();
   rows.pop();
+  function getValue(row: number, column: number) {
+    return getCell({ rows, columns }, row, column);
+  }
+  // check whether there is actual data in the table
+  if (
+    rows.every((_, index) => !getValue(index, metricIndex) && !getValue(index, metricIndex + 1))
+  ) {
+    throw new Error('all cell contents falsy');
+  }
   rows.forEach((_, index) => {
     if (index < rows.length - 1) {
       expect(getCell({ rows, columns }, index, metricIndex + 1)).to.be(
@@ -137,6 +146,22 @@ export default function ({
       checkShift(result.rows, result.columns);
     });
 
+    it('shifts correctly even if one id is the prefix of another', async () => {
+      const expression = `
+          kibana_context timeRange={timerange from='${timeRange.from}' to='${timeRange.to}'}
+          | esaggs index={indexPatternLoad id='logstash-*'}
+          aggs={aggDateHistogram id="prefix" enabled=true schema="bucket" field="@timestamp" interval="1h"}
+          aggs={aggAvg id="prefix-prefix" field="bytes" enabled=true schema="metric" timeShift="1h"}
+          aggs={aggAvg id="prefix-prefix-prefix" field="bytes" enabled=true schema="metric"}
+        `;
+      const result: Datatable = await expectExpression(
+        'esaggs_shift_date_histogram',
+        expression
+      ).getResponse();
+      expect(result.rows.length).to.be(25);
+      checkShift(result.rows, result.columns);
+    });
+
     it('shifts filtered metrics', async () => {
       const expression = `
           kibana_context timeRange={timerange from='${timeRange.from}' to='${timeRange.to}'}