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

import type {
  QueryDSL,
  BoolQuery,
  MatchQuery,
  MatchPhraseQuery,
  MultiMatchQuery,
  TermQuery,
  TermsQuery,
  RangeQuery,
  ExistsQuery,
  PrefixQuery,
  WildcardQuery,
  RegexpQuery,
  FuzzyQuery,
  IdsQuery,
  MatchAllQuery,
  QueryStringQuery,
  SimpleQueryStringQuery,
  SearchOptions,
  SearchResult,
  SortOrder,
  MatchOperator,
  MultiMatchType,
  RangeBounds,
  SortField,
} from './types.js';
import type { AggregationBuilder } from './aggregation-builder.js';
import { createAggregationBuilder } from './aggregation-builder.js';
import { ElasticsearchConnectionManager } from '../../core/connection/connection-manager.js';
import { defaultLogger } from '../../core/observability/logger.js';
import { defaultMetrics } from '../../core/observability/metrics.js';
import { defaultTracing } from '../../core/observability/tracing.js';

/**
 * Fluent query builder for type-safe Elasticsearch queries
 *
 * @example
 * ```typescript
 * const results = await new QueryBuilder<Product>('products')
 *   .match('name', 'laptop')
 *   .range('price', { gte: 100, lte: 1000 })
 *   .sort('price', 'asc')
 *   .size(20)
 *   .execute();
 * ```
 */
export class QueryBuilder<T = unknown> {
  private index: string;
  private queryDSL: QueryDSL | null = null;
  private boolClauses: {
    must: QueryDSL[];
    should: QueryDSL[];
    must_not: QueryDSL[];
    filter: QueryDSL[];
  } = {
    must: [],
    should: [],
    must_not: [],
    filter: [],
  };
  private minimumShouldMatch?: number | string;
  private sortFields: Array<SortField | string> = [];
  private sourceFields?: string[];
  private excludeSourceFields?: string[];
  private resultSize: number = 10;
  private resultFrom: number = 0;
  private shouldTrackTotalHits: boolean | number = true;
  private searchTimeout?: string;
  private aggregationBuilder?: AggregationBuilder;
  private highlightConfig?: SearchOptions['highlight'];

  constructor(index: string) {
    this.index = index;
  }

  /**
   * Create a new query builder instance
   */
  static create<T>(index: string): QueryBuilder<T> {
    return new QueryBuilder<T>(index);
  }

  // ============================================================================
  // Query Methods
  // ============================================================================

  /**
   * Add a match query
   */
  match(field: string, query: string, options?: { operator?: MatchOperator; fuzziness?: number | 'AUTO'; boost?: number }): this {
    const matchQuery: MatchQuery = {
      match: {
        [field]: {
          query,
          ...options,
        },
      },
    };
    this.boolClauses.must.push(matchQuery);
    return this;
  }

  /**
   * Add a match phrase query
   */
  matchPhrase(field: string, query: string, options?: { slop?: number; boost?: number }): this {
    const matchPhraseQuery: MatchPhraseQuery = {
      match_phrase: {
        [field]: {
          query,
          ...options,
        },
      },
    };
    this.boolClauses.must.push(matchPhraseQuery);
    return this;
  }

  /**
   * Add a multi-match query
   */
  multiMatch(query: string, fields: string[], options?: { type?: MultiMatchType; operator?: MatchOperator; boost?: number }): this {
    const multiMatchQuery: MultiMatchQuery = {
      multi_match: {
        query,
        fields,
        ...options,
      },
    };
    this.boolClauses.must.push(multiMatchQuery);
    return this;
  }

  /**
   * Add a term query (exact match)
   */
  term(field: string, value: string | number | boolean, boost?: number): this {
    const termQuery: TermQuery = {
      term: {
        [field]: {
          value,
          ...(boost && { boost }),
        },
      },
    };
    this.boolClauses.filter.push(termQuery);
    return this;
  }

  /**
   * Add a terms query (match any of the values)
   */
  terms(field: string, values: Array<string | number | boolean>, boost?: number): this {
    const termsQuery: TermsQuery = {
      terms: {
        [field]: values,
        ...(boost && { boost }),
      },
    };
    this.boolClauses.filter.push(termsQuery);
    return this;
  }

  /**
   * Add a range query
   */
  range(field: string, bounds: RangeBounds, boost?: number): this {
    const rangeQuery: RangeQuery = {
      range: {
        [field]: {
          ...bounds,
          ...(boost && { boost }),
        },
      },
    };
    this.boolClauses.filter.push(rangeQuery);
    return this;
  }

  /**
   * Add an exists query (field must exist)
   */
  exists(field: string): this {
    const existsQuery: ExistsQuery = {
      exists: {
        field,
      },
    };
    this.boolClauses.filter.push(existsQuery);
    return this;
  }

  /**
   * Add a prefix query
   */
  prefix(field: string, value: string, boost?: number): this {
    const prefixQuery: PrefixQuery = {
      prefix: {
        [field]: {
          value,
          ...(boost && { boost }),
        },
      },
    };
    this.boolClauses.must.push(prefixQuery);
    return this;
  }

  /**
   * Add a wildcard query
   */
  wildcard(field: string, value: string, boost?: number): this {
    const wildcardQuery: WildcardQuery = {
      wildcard: {
        [field]: {
          value,
          ...(boost && { boost }),
        },
      },
    };
    this.boolClauses.must.push(wildcardQuery);
    return this;
  }

  /**
   * Add a regexp query
   */
  regexp(field: string, value: string, options?: { flags?: string; boost?: number }): this {
    const regexpQuery: RegexpQuery = {
      regexp: {
        [field]: {
          value,
          ...options,
        },
      },
    };
    this.boolClauses.must.push(regexpQuery);
    return this;
  }

  /**
   * Add a fuzzy query
   */
  fuzzy(field: string, value: string, options?: { fuzziness?: number | 'AUTO'; boost?: number }): this {
    const fuzzyQuery: FuzzyQuery = {
      fuzzy: {
        [field]: {
          value,
          ...options,
        },
      },
    };
    this.boolClauses.must.push(fuzzyQuery);
    return this;
  }

  /**
   * Add an IDs query
   */
  ids(values: string[]): this {
    const idsQuery: IdsQuery = {
      ids: {
        values,
      },
    };
    this.boolClauses.filter.push(idsQuery);
    return this;
  }

  /**
   * Add a query string query
   */
  queryString(query: string, options?: { default_field?: string; fields?: string[]; default_operator?: MatchOperator; boost?: number }): this {
    const queryStringQuery: QueryStringQuery = {
      query_string: {
        query,
        ...options,
      },
    };
    this.boolClauses.must.push(queryStringQuery);
    return this;
  }

  /**
   * Add a simple query string query
   */
  simpleQueryString(query: string, options?: { fields?: string[]; default_operator?: MatchOperator; boost?: number }): this {
    const simpleQueryStringQuery: SimpleQueryStringQuery = {
      simple_query_string: {
        query,
        ...options,
      },
    };
    this.boolClauses.must.push(simpleQueryStringQuery);
    return this;
  }

  /**
   * Match all documents
   */
  matchAll(boost?: number): this {
    const matchAllQuery: MatchAllQuery = {
      match_all: {
        ...(boost && { boost }),
      },
    };
    this.queryDSL = matchAllQuery;
    return this;
  }

  /**
   * Add a custom query to the must clause
   */
  must(query: QueryDSL): this {
    this.boolClauses.must.push(query);
    return this;
  }

  /**
   * Add a custom query to the should clause
   */
  should(query: QueryDSL): this {
    this.boolClauses.should.push(query);
    return this;
  }

  /**
   * Add a custom query to the must_not clause
   */
  mustNot(query: QueryDSL): this {
    this.boolClauses.must_not.push(query);
    return this;
  }

  /**
   * Add a custom query to the filter clause
   */
  filter(query: QueryDSL): this {
    this.boolClauses.filter.push(query);
    return this;
  }

  /**
   * Set minimum_should_match for boolean queries
   */
  minimumMatch(value: number | string): this {
    this.minimumShouldMatch = value;
    return this;
  }

  /**
   * Set a custom query DSL (replaces builder queries)
   */
  customQuery(query: QueryDSL): this {
    this.queryDSL = query;
    return this;
  }

  // ============================================================================
  // Result Shaping Methods
  // ============================================================================

  /**
   * Add sorting
   */
  sort(field: string, order: SortOrder = 'asc'): this {
    this.sortFields.push({ [field]: { order } });
    return this;
  }

  /**
   * Add custom sort configuration
   */
  customSort(sort: SortField | string): this {
    this.sortFields.push(sort);
    return this;
  }

  /**
   * Specify fields to include in results (source filtering)
   */
  fields(fields: string[]): this {
    this.sourceFields = fields;
    return this;
  }

  /**
   * Specify fields to exclude from results
   */
  exclude(fields: string[]): this {
    this.excludeSourceFields = fields;
    return this;
  }

  /**
   * Set number of results to return
   */
  size(size: number): this {
    this.resultSize = size;
    return this;
  }

  /**
   * Set offset for pagination
   */
  from(from: number): this {
    this.resultFrom = from;
    return this;
  }

  /**
   * Set pagination (convenience method)
   */
  paginate(page: number, pageSize: number): this {
    this.resultFrom = (page - 1) * pageSize;
    this.resultSize = pageSize;
    return this;
  }

  /**
   * Set whether to track total hits
   */
  trackTotalHits(track: boolean | number): this {
    this.shouldTrackTotalHits = track;
    return this;
  }

  /**
   * Set search timeout
   */
  timeout(timeout: string): this {
    this.searchTimeout = timeout;
    return this;
  }

  /**
   * Configure highlighting
   */
  highlight(config: SearchOptions['highlight']): this {
    this.highlightConfig = config;
    return this;
  }

  // ============================================================================
  // Aggregation Methods
  // ============================================================================

  /**
   * Get aggregation builder
   */
  aggregations(configure: (builder: AggregationBuilder) => void): this {
    if (!this.aggregationBuilder) {
      this.aggregationBuilder = createAggregationBuilder();
    }
    configure(this.aggregationBuilder);
    return this;
  }

  // ============================================================================
  // Build & Execute
  // ============================================================================

  /**
   * Build the final query DSL
   */
  build(): SearchOptions {
    let finalQuery: QueryDSL | undefined;

    // If custom query was set, use it
    if (this.queryDSL) {
      finalQuery = this.queryDSL;
    } else {
      // Otherwise, build from bool clauses
      const hasAnyClauses =
        this.boolClauses.must.length > 0 ||
        this.boolClauses.should.length > 0 ||
        this.boolClauses.must_not.length > 0 ||
        this.boolClauses.filter.length > 0;

      if (hasAnyClauses) {
        const boolQuery: BoolQuery = {
          bool: {},
        };

        if (this.boolClauses.must.length > 0) {
          boolQuery.bool.must = this.boolClauses.must;
        }
        if (this.boolClauses.should.length > 0) {
          boolQuery.bool.should = this.boolClauses.should;
        }
        if (this.boolClauses.must_not.length > 0) {
          boolQuery.bool.must_not = this.boolClauses.must_not;
        }
        if (this.boolClauses.filter.length > 0) {
          boolQuery.bool.filter = this.boolClauses.filter;
        }
        if (this.minimumShouldMatch !== undefined) {
          boolQuery.bool.minimum_should_match = this.minimumShouldMatch;
        }

        finalQuery = boolQuery;
      }
    }

    const searchOptions: SearchOptions = {
      ...(finalQuery && { query: finalQuery }),
      ...(this.sourceFields && { fields: this.sourceFields }),
      ...(this.excludeSourceFields && { excludeFields: this.excludeSourceFields }),
      size: this.resultSize,
      from: this.resultFrom,
      ...(this.sortFields.length > 0 && { sort: this.sortFields }),
      trackTotalHits: this.shouldTrackTotalHits,
      ...(this.searchTimeout && { timeout: this.searchTimeout }),
      ...(this.highlightConfig && { highlight: this.highlightConfig }),
      ...(this.aggregationBuilder && { aggregations: this.aggregationBuilder.build() }),
    };

    return searchOptions;
  }

  /**
   * Execute the query and return results
   */
  async execute(): Promise<SearchResult<T>> {
    const span = defaultTracing.createSpan('query.execute', {
      'db.system': 'elasticsearch',
      'db.operation': 'search',
      'db.elasticsearch.index': this.index,
    });

    try {
      const client = ElasticsearchConnectionManager.getInstance().getClient();
      const searchOptions = this.build();

      defaultLogger.debug('Executing query', {
        index: this.index,
        query: searchOptions.query,
        size: searchOptions.size,
        from: searchOptions.from,
      });

      const startTime = Date.now();

      // Execute search
      const result = await client.search<T>({
        index: this.index,
        ...searchOptions,
      });

      const duration = Date.now() - startTime;

      // Record metrics
      defaultMetrics.requestsTotal.inc({ operation: 'search', index: this.index });
      defaultMetrics.requestDuration.observe({ operation: 'search', index: this.index }, duration);

      defaultLogger.info('Query executed successfully', {
        index: this.index,
        took: result.took,
        hits: result.hits.total,
        duration,
      });

      span.setAttributes({
        'db.elasticsearch.took': result.took,
        'db.elasticsearch.hits': typeof result.hits.total === 'object' ? result.hits.total.value : result.hits.total,
      });
      span.end();

      return result as SearchResult<T>;
    } catch (error) {
      defaultMetrics.requestErrors.inc({ operation: 'search', index: this.index });
      defaultLogger.error('Query execution failed', { index: this.index, error: error instanceof Error ? error.message : String(error) });
      span.recordException(error as Error);
      span.end();
      throw error;
    }
  }

  /**
   * Execute query and return only the hits
   */
  async executeAndGetHits(): Promise<SearchResult<T>['hits']['hits']> {
    const result = await this.execute();
    return result.hits.hits;
  }

  /**
   * Execute query and return only the source documents
   */
  async executeAndGetSources(): Promise<T[]> {
    const hits = await this.executeAndGetHits();
    return hits.map((hit) => hit._source);
  }

  /**
   * Count documents matching the query
   */
  async count(): Promise<number> {
    const span = defaultTracing.createSpan('query.count', {
      'db.system': 'elasticsearch',
      'db.operation': 'count',
      'db.elasticsearch.index': this.index,
    });

    try {
      const client = ElasticsearchConnectionManager.getInstance().getClient();
      const searchOptions = this.build();

      const result = await client.count({
        index: this.index,
        ...(searchOptions.query && { query: searchOptions.query }),
      });

      span.end();
      return result.count;
    } catch (error) {
      span.recordException(error as Error);
      span.end();
      throw error;
    }
  }
}

/**
 * Create a new query builder instance
 */
export function createQuery<T>(index: string): QueryBuilder<T> {
  return new QueryBuilder<T>(index);
}