elasticsearch/ts/domain/query/aggregation-builder.ts

import type {
  AggregationDSL,
  TermsAggregation,
  MetricAggregation,
  StatsAggregation,
  ExtendedStatsAggregation,
  PercentilesAggregation,
  DateHistogramAggregation,
  HistogramAggregation,
  RangeAggregation,
  FilterAggregation,
  TopHitsAggregation,
  QueryDSL,
  SortOrder,
  SortField,
} from './types.js';

/**
 * Fluent aggregation builder for type-safe Elasticsearch aggregations
 *
 * @example
 * ```typescript
 * const query = new QueryBuilder<Product>('products')
 *   .aggregations((agg) => {
 *     agg.terms('categories', 'category.keyword', { size: 10 })
 *       .subAggregation('avg_price', (sub) => sub.avg('price'));
 *   });
 * ```
 */
export class AggregationBuilder {
  private aggregations: Record<string, AggregationDSL> = {};
  private currentAggName?: string;

  /**
   * Add a terms aggregation
   */
  terms(
    name: string,
    field: string,
    options?: {
      size?: number;
      order?: Record<string, SortOrder>;
      missing?: string | number;
    }
  ): this {
    const termsAgg: TermsAggregation = {
      terms: {
        field,
        ...options,
      },
    };
    this.aggregations[name] = termsAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add an average metric aggregation
   */
  avg(name: string, field: string, missing?: number): this {
    const avgAgg: MetricAggregation = {
      avg: {
        field,
        ...(missing !== undefined && { missing }),
      },
    };
    this.aggregations[name] = avgAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a sum metric aggregation
   */
  sum(name: string, field: string, missing?: number): this {
    const sumAgg: MetricAggregation = {
      sum: {
        field,
        ...(missing !== undefined && { missing }),
      },
    };
    this.aggregations[name] = sumAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a min metric aggregation
   */
  min(name: string, field: string, missing?: number): this {
    const minAgg: MetricAggregation = {
      min: {
        field,
        ...(missing !== undefined && { missing }),
      },
    };
    this.aggregations[name] = minAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a max metric aggregation
   */
  max(name: string, field: string, missing?: number): this {
    const maxAgg: MetricAggregation = {
      max: {
        field,
        ...(missing !== undefined && { missing }),
      },
    };
    this.aggregations[name] = maxAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a cardinality metric aggregation
   */
  cardinality(name: string, field: string): this {
    const cardinalityAgg: MetricAggregation = {
      cardinality: {
        field,
      },
    };
    this.aggregations[name] = cardinalityAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a stats aggregation
   */
  stats(name: string, field: string): this {
    const statsAgg: StatsAggregation = {
      stats: {
        field,
      },
    };
    this.aggregations[name] = statsAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add an extended stats aggregation
   */
  extendedStats(name: string, field: string): this {
    const extendedStatsAgg: ExtendedStatsAggregation = {
      extended_stats: {
        field,
      },
    };
    this.aggregations[name] = extendedStatsAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a percentiles aggregation
   */
  percentiles(name: string, field: string, percents?: number[]): this {
    const percentilesAgg: PercentilesAggregation = {
      percentiles: {
        field,
        ...(percents && { percents }),
      },
    };
    this.aggregations[name] = percentilesAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a date histogram aggregation
   */
  dateHistogram(
    name: string,
    field: string,
    options: {
      calendar_interval?: string;
      fixed_interval?: string;
      format?: string;
      time_zone?: string;
      min_doc_count?: number;
    }
  ): this {
    const dateHistogramAgg: DateHistogramAggregation = {
      date_histogram: {
        field,
        ...options,
      },
    };
    this.aggregations[name] = dateHistogramAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a histogram aggregation
   */
  histogram(
    name: string,
    field: string,
    interval: number,
    options?: {
      min_doc_count?: number;
    }
  ): this {
    const histogramAgg: HistogramAggregation = {
      histogram: {
        field,
        interval,
        ...options,
      },
    };
    this.aggregations[name] = histogramAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a range aggregation
   */
  range(
    name: string,
    field: string,
    ranges: Array<{ from?: number; to?: number; key?: string }>
  ): this {
    const rangeAgg: RangeAggregation = {
      range: {
        field,
        ranges,
      },
    };
    this.aggregations[name] = rangeAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a filter aggregation
   */
  filterAgg(name: string, filter: QueryDSL): this {
    const filterAgg: FilterAggregation = {
      filter,
    };
    this.aggregations[name] = filterAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a top hits aggregation
   */
  topHits(
    name: string,
    options?: {
      size?: number;
      sort?: Array<SortField | string>;
      _source?: boolean | { includes?: string[]; excludes?: string[] };
    }
  ): this {
    const topHitsAgg: TopHitsAggregation = {
      top_hits: {
        ...options,
      },
    };
    this.aggregations[name] = topHitsAgg;
    this.currentAggName = name;
    return this;
  }

  /**
   * Add a sub-aggregation to the last defined aggregation
   */
  subAggregation(name: string, configure: (builder: AggregationBuilder) => void): this {
    if (!this.currentAggName) {
      throw new Error('Cannot add sub-aggregation: no parent aggregation defined');
    }

    const parentAgg = this.aggregations[this.currentAggName];
    const subBuilder = new AggregationBuilder();
    configure(subBuilder);

    // Add aggs field to parent aggregation
    if ('terms' in parentAgg) {
      (parentAgg as TermsAggregation).aggs = subBuilder.build();
    } else if ('date_histogram' in parentAgg) {
      (parentAgg as DateHistogramAggregation).aggs = subBuilder.build();
    } else if ('histogram' in parentAgg) {
      (parentAgg as HistogramAggregation).aggs = subBuilder.build();
    } else if ('range' in parentAgg) {
      (parentAgg as RangeAggregation).aggs = subBuilder.build();
    } else if ('filter' in parentAgg) {
      (parentAgg as FilterAggregation).aggs = subBuilder.build();
    }

    return this;
  }

  /**
   * Add a custom aggregation DSL
   */
  custom(name: string, aggregation: AggregationDSL): this {
    this.aggregations[name] = aggregation;
    this.currentAggName = name;
    return this;
  }

  /**
   * Build the aggregations object
   */
  build(): Record<string, AggregationDSL> {
    return this.aggregations;
  }
}

/**
 * Create a new aggregation builder
 */
export function createAggregationBuilder(): AggregationBuilder {
  return new AggregationBuilder();
}