code: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   *
   */
  
  package org.apache.skywalking.oap.server.storage.plugin.elasticsearch.query;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import org.apache.skywalking.library.elasticsearch.requests.search.BoolQueryBuilder;
  import org.apache.skywalking.library.elasticsearch.requests.search.Query;
  import org.apache.skywalking.library.elasticsearch.requests.search.Search;
  import org.apache.skywalking.library.elasticsearch.requests.search.SearchBuilder;
  import org.apache.skywalking.library.elasticsearch.requests.search.aggregation.Aggregation;
  import org.apache.skywalking.library.elasticsearch.requests.search.aggregation.TermsAggregationBuilder;
  import org.apache.skywalking.library.elasticsearch.response.search.SearchResponse;
  import org.apache.skywalking.oap.server.core.UnexpectedException;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.endpoint.EndpointRelationServerSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.instance.ServiceInstanceRelationClientSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.instance.ServiceInstanceRelationServerSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.process.ProcessRelationClientSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.process.ProcessRelationServerSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.service.ServiceRelationClientSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.manual.relation.service.ServiceRelationServerSideMetrics;
  import org.apache.skywalking.oap.server.core.analysis.metrics.IntList;
  import org.apache.skywalking.oap.server.core.analysis.metrics.Metrics;
  import org.apache.skywalking.oap.server.core.query.input.Duration;
  import org.apache.skywalking.oap.server.core.query.type.Call;
  import org.apache.skywalking.oap.server.core.source.DetectPoint;
  import org.apache.skywalking.oap.server.core.storage.query.ITopologyQueryDAO;
  import org.apache.skywalking.oap.server.library.client.elasticsearch.ElasticSearchClient;
  import org.apache.skywalking.oap.server.library.util.CollectionUtils;
  import org.apache.skywalking.oap.server.storage.plugin.elasticsearch.base.EsDAO;
  import org.apache.skywalking.oap.server.storage.plugin.elasticsearch.base.IndexController;
  
  public class TopologyQueryEsDAO extends EsDAO implements ITopologyQueryDAO {
  
      public TopologyQueryEsDAO(ElasticSearchClient client) {
          super(client);
      }
  
      @Override
      public List<Call.CallDetail> loadServiceRelationsDetectedAtServerSide(Duration duration, List<String> serviceIds) {
          if (CollectionUtils.isEmpty(serviceIds)) {
              throw new UnexpectedException("Service id is empty");
          }
  
          final SearchBuilder sourceBuilder = Search.builder().size(0);
          setQueryCondition(sourceBuilder, duration, serviceIds, ServiceRelationServerSideMetrics.INDEX_NAME);
  
          return buildServiceRelation(
              sourceBuilder, ServiceRelationServerSideMetrics.INDEX_NAME, DetectPoint.SERVER);
      }
  
      @Override
      public List<Call.CallDetail> loadServiceRelationDetectedAtClientSide(Duration duration,
                                                                           List<String> serviceIds) {
          if (CollectionUtils.isEmpty(serviceIds)) {
              throw new UnexpectedException("Service id is empty");
          }
  
          final SearchBuilder sourceBuilder = Search.builder().size(0);
          setQueryCondition(sourceBuilder, duration, serviceIds, ServiceRelationClientSideMetrics.INDEX_NAME);
  
          return buildServiceRelation(
              sourceBuilder, ServiceRelationClientSideMetrics.INDEX_NAME, DetectPoint.CLIENT);
      }
  
      @Override
      public List<Call.CallDetail> loadServiceRelationsDetectedAtServerSide(Duration duration) {
          SearchBuilder sourceBuilder = Search.builder();
          final BoolQueryBuilder query = Query.bool()
                                              .must(Query.range(ServiceRelationServerSideMetrics.TIME_BUCKET)
                                                         .gte(duration.getStartTimeBucket())
                                                         .lte(duration.getEndTimeBucket()));
          if (IndexController.LogicIndicesRegister.isMergedTable(ServiceRelationServerSideMetrics.INDEX_NAME)) {
              query.must(Query.term(IndexController.LogicIndicesRegister.METRIC_TABLE_NAME,
                                    ServiceRelationServerSideMetrics.INDEX_NAME
              ));
          }
          sourceBuilder.query(query).size(0);
  
          return buildServiceRelation(
              sourceBuilder, ServiceRelationServerSideMetrics.INDEX_NAME, DetectPoint.SERVER);
      }
  
      @Override
      public List<Call.CallDetail> loadServiceRelationDetectedAtClientSide(Duration duration) {
          SearchBuilder sourceBuilder = Search.builder();
          final BoolQueryBuilder query = Query.bool()
                                              .must(Query.range(ServiceRelationClientSideMetrics.TIME_BUCKET)
                                                         .gte(duration.getStartTimeBucket())
                                                         .lte(duration.getEndTimeBucket()));
          if (IndexController.LogicIndicesRegister.isMergedTable(ServiceRelationClientSideMetrics.INDEX_NAME)) {
              query.must(Query.term(IndexController.LogicIndicesRegister.METRIC_TABLE_NAME,
                                    ServiceRelationClientSideMetrics.INDEX_NAME
              ));
          }
          sourceBuilder.query(query).size(0);
  
          return buildServiceRelation(
              sourceBuilder, ServiceRelationClientSideMetrics.INDEX_NAME, DetectPoint.CLIENT);
      }
  
      @Override
      public List<Call.CallDetail> loadInstanceRelationDetectedAtServerSide(String clientServiceId,
                                                                            String serverServiceId,
                                                                            Duration duration) {
          final SearchBuilder search = Search.builder().size(0);
          setInstanceQueryCondition(
              search, duration, clientServiceId, serverServiceId, ServiceInstanceRelationServerSideMetrics.INDEX_NAME);
  
          return buildInstanceRelation(
              search, ServiceInstanceRelationServerSideMetrics.INDEX_NAME, DetectPoint.SERVER);
      }
  
      @Override
      public List<Call.CallDetail> loadInstanceRelationDetectedAtClientSide(String clientServiceId,
                                                                            String serverServiceId,
                                                                            Duration duration) {
          final SearchBuilder search = Search.builder().size(0);
          setInstanceQueryCondition(
              search, duration, clientServiceId, serverServiceId, ServiceInstanceRelationClientSideMetrics.INDEX_NAME);
  
          return buildInstanceRelation(
              search, ServiceInstanceRelationClientSideMetrics.INDEX_NAME, DetectPoint.CLIENT);
      }
  
      private void setInstanceQueryCondition(SearchBuilder search, Duration duration,
                                             String clientServiceId, String serverServiceId, String indexName) {
          final BoolQueryBuilder serverRelationBoolQuery =
              Query.bool()
                   .must(
                       Query.term(
                           ServiceInstanceRelationServerSideMetrics.SOURCE_SERVICE_ID,
                           clientServiceId
                       ))
                   .must(
                       Query.term(
                           ServiceInstanceRelationServerSideMetrics.DEST_SERVICE_ID,
                           serverServiceId
                       ));
  
          final BoolQueryBuilder clientRelationBoolQuery =
              Query.bool()
                   .must(
                       Query.term(
                           ServiceInstanceRelationServerSideMetrics.DEST_SERVICE_ID,
                           clientServiceId
                       ))
                   .must(
                       Query.term(
                           ServiceInstanceRelationServerSideMetrics.SOURCE_SERVICE_ID,
                           serverServiceId
                       ));
  
          final BoolQueryBuilder serviceIdBoolQuery =
              Query.bool()
                   .should(serverRelationBoolQuery)
                   .should(clientRelationBoolQuery);
  
          final BoolQueryBuilder boolQuery =
              Query.bool()
                   .must(Query.range(EndpointRelationServerSideMetrics.TIME_BUCKET)
                              .gte(duration.getStartTimeBucket())
                              .lte(duration.getEndTimeBucket()))
                   .must(serviceIdBoolQuery);
          if (IndexController.LogicIndicesRegister.isMergedTable(indexName)) {
              boolQuery.must(Query.term(IndexController.LogicIndicesRegister.METRIC_TABLE_NAME, indexName));
          }
          search.query(boolQuery);
      }
  
      @Override
      public List<Call.CallDetail> loadEndpointRelation(Duration duration, String destEndpointId) {
          SearchBuilder sourceBuilder = Search.builder();
          sourceBuilder.size(0);
  
          BoolQueryBuilder boolQuery = Query.bool();
          boolQuery.must(
              Query.range(EndpointRelationServerSideMetrics.TIME_BUCKET)
                   .gte(duration.getStartTimeBucket()).lte(duration.getEndTimeBucket()));
  
          BoolQueryBuilder serviceIdBoolQuery = Query.bool();
          boolQuery.must(serviceIdBoolQuery);
          serviceIdBoolQuery.should(
              Query.term(
                  EndpointRelationServerSideMetrics.SOURCE_ENDPOINT, destEndpointId
              ));
          serviceIdBoolQuery.should(
              Query.term(
                  EndpointRelationServerSideMetrics.DEST_ENDPOINT, destEndpointId
              ));
          if (IndexController.LogicIndicesRegister.isMergedTable(EndpointRelationServerSideMetrics.INDEX_NAME)) {
              boolQuery.must(Query.term(IndexController.LogicIndicesRegister.METRIC_TABLE_NAME,
                                        EndpointRelationServerSideMetrics.INDEX_NAME
              ));
          }
          sourceBuilder.query(boolQuery);
  
          return loadEndpoint(
              sourceBuilder, EndpointRelationServerSideMetrics.INDEX_NAME, DetectPoint.SERVER);
      }
  
      @Override
      public List<Call.CallDetail> loadProcessRelationDetectedAtClientSide(String serviceInstanceId,
                                                                           Duration duration) throws IOException {
          return buildProcessRelation(serviceInstanceId, duration, DetectPoint.CLIENT);
      }
  
      @Override
      public List<Call.CallDetail> loadProcessRelationDetectedAtServerSide(String serviceInstanceId,
                                                                           Duration duration) throws IOException {
          return buildProcessRelation(serviceInstanceId, duration, DetectPoint.SERVER);
      }
  
      private List<Call.CallDetail> buildProcessRelation(String serviceInstanceId,
                                                         Duration duration,
                                                         DetectPoint detectPoint) throws IOException {
          final SearchBuilder sourceBuilder = Search.builder().size(0);
          final BoolQueryBuilder query = Query.bool()
                                              .must(Query.term(ProcessRelationServerSideMetrics.SERVICE_INSTANCE_ID,
                                                               serviceInstanceId
                                              ))
                                              .must(Query.range(EndpointRelationServerSideMetrics.TIME_BUCKET)
                                                         .gte(duration.getStartTimeBucket())
                                                         .lte(duration.getEndTimeBucket()));
          sourceBuilder.query(query);
          sourceBuilder.aggregation(
              Aggregation
                  .terms(Metrics.ENTITY_ID).field(Metrics.ENTITY_ID)
                  .subAggregation(
                      Aggregation.terms(ProcessRelationServerSideMetrics.COMPONENT_ID)
                                 .field(ProcessRelationServerSideMetrics.COMPONENT_ID)
                                 .executionHint(TermsAggregationBuilder.ExecutionHint.MAP)
                                 .collectMode(TermsAggregationBuilder.CollectMode.BREADTH_FIRST))
                  .executionHint(TermsAggregationBuilder.ExecutionHint.MAP)
                  .collectMode(TermsAggregationBuilder.CollectMode.BREADTH_FIRST)
                  .size(1000));
  
          String indexName = detectPoint == DetectPoint.SERVER ?
              ProcessRelationServerSideMetrics.INDEX_NAME : ProcessRelationClientSideMetrics.INDEX_NAME;
  
          final String index = IndexController.LogicIndicesRegister.getPhysicalTableName(indexName);
          if (IndexController.LogicIndicesRegister.isMergedTable(indexName)) {
              query.must(Query.term(IndexController.LogicIndicesRegister.METRIC_TABLE_NAME, indexName));
          }
          final SearchResponse response = getClient().search(index, sourceBuilder.build());
  
          final List<Call.CallDetail> calls = new ArrayList<>();
          final Map<String, Object> entityTerms =
              (Map<String, Object>) response.getAggregations().get(Metrics.ENTITY_ID);
          final List<Map<String, Object>> buckets =
              (List<Map<String, Object>>) entityTerms.get("buckets");
          for (final Map<String, Object> entityBucket : buckets) {
              String entityId = (String) entityBucket.get("key");
              final Map<String, Object> componentTerms =
                  (Map<String, Object>) entityBucket.get(
                      ProcessRelationServerSideMetrics.COMPONENT_ID);
              final List<Map<String, Object>> subAgg =
                  (List<Map<String, Object>>) componentTerms.get("buckets");
              final int componentId = ((Number) subAgg.iterator().next().get("key")).intValue();
  
              Call.CallDetail call = new Call.CallDetail();
              call.buildProcessRelation(entityId, componentId, detectPoint);
              calls.add(call);
          }
          return calls;
      }
  
      private List<Call.CallDetail> buildServiceRelation(SearchBuilder sourceBuilder,
                                                         String indexName,
                                                         DetectPoint detectPoint) {
          sourceBuilder.aggregation(
              Aggregation
                  .terms(Metrics.ENTITY_ID).field(Metrics.ENTITY_ID)
                  .subAggregation(
                      Aggregation.terms(ServiceRelationServerSideMetrics.COMPONENT_IDS)
                                 .field(ServiceRelationServerSideMetrics.COMPONENT_IDS)
                                 .executionHint(TermsAggregationBuilder.ExecutionHint.MAP)
                                 .collectMode(TermsAggregationBuilder.CollectMode.BREADTH_FIRST))
                  .executionHint(TermsAggregationBuilder.ExecutionHint.MAP)
                  .collectMode(TermsAggregationBuilder.CollectMode.BREADTH_FIRST)
                  .size(1000));
  
          final String index =
              IndexController.LogicIndicesRegister.getPhysicalTableName(indexName);
          final SearchResponse response = getClient().search(index, sourceBuilder.build());
  
          final List<Call.CallDetail> calls = new ArrayList<>();
          final Map<String, Object> entityTerms =
              (Map<String, Object>) response.getAggregations().get(Metrics.ENTITY_ID);
          final List<Map<String, Object>> buckets =
              (List<Map<String, Object>>) entityTerms.get("buckets");
          for (final Map<String, Object> entityBucket : buckets) {
              String entityId = (String) entityBucket.get("key");
              final Map<String, Object> componentTerms =
                  (Map<String, Object>) entityBucket.get(
                      ServiceRelationServerSideMetrics.COMPONENT_IDS);
              final List<Map<String, Object>> subAgg =
                  (List<Map<String, Object>>) componentTerms.get("buckets");
              final IntList componentIds = new IntList((String) subAgg.iterator().next().get("key"));
  
              for (int i = 0; i < componentIds.size(); i++) {
                  Call.CallDetail call = new Call.CallDetail();
                  call.buildFromServiceRelation(entityId, componentIds.get(i), detectPoint);
                  calls.add(call);
              }
          }
          return calls;
      }
  
      private List<Call.CallDetail> buildInstanceRelation(SearchBuilder sourceBuilder,
                                                          String indexName,
                                                          DetectPoint detectPoint) {
          sourceBuilder.aggregation(
              Aggregation
                  .terms(Metrics.ENTITY_ID).field(Metrics.ENTITY_ID)
                  .executionHint(TermsAggregationBuilder.ExecutionHint.MAP)
                  .collectMode(TermsAggregationBuilder.CollectMode.BREADTH_FIRST)
                  .size(1000));
  
          final String index =
              IndexController.LogicIndicesRegister.getPhysicalTableName(indexName);
          SearchResponse response = getClient().search(index, sourceBuilder.build());
  
          List<Call.CallDetail> calls = new ArrayList<>();
          final Map<String, Object> entityTerms =
              (Map<String, Object>) response.getAggregations().get(Metrics.ENTITY_ID);
  
          final List<Map<String, Object>> buckets =
              (List<Map<String, Object>>) entityTerms.get("buckets");
          for (final Map<String, Object> entityBucket : buckets) {
              String entityId = (String) entityBucket.get("key");
  
              Call.CallDetail call = new Call.CallDetail();
              call.buildFromInstanceRelation(entityId, detectPoint);
              calls.add(call);
          }
          return calls;
      }
  
      private List<Call.CallDetail> loadEndpoint(SearchBuilder sourceBuilder, String indexName,
                                                 DetectPoint detectPoint) {
          sourceBuilder.aggregation(
              Aggregation.terms(Metrics.ENTITY_ID).field(Metrics.ENTITY_ID)
                         .executionHint(TermsAggregationBuilder.ExecutionHint.MAP)
                         .collectMode(TermsAggregationBuilder.CollectMode.BREADTH_FIRST)
                         .size(1000));
  
          final String index =
              IndexController.LogicIndicesRegister.getPhysicalTableName(indexName);
          final SearchResponse response = getClient().search(index, sourceBuilder.build());
  
          final List<Call.CallDetail> calls = new ArrayList<>();
          final Map<String, Object> entityTerms =
              (Map<String, Object>) response.getAggregations().get(Metrics.ENTITY_ID);
          final List<Map<String, Object>> buckets =
              (List<Map<String, Object>>) entityTerms.get("buckets");
          for (final Map<String, Object> entityBucket : buckets) {
              String entityId = (String) entityBucket.get("key");
  
              Call.CallDetail call = new Call.CallDetail();
              call.buildFromEndpointRelation(entityId, detectPoint);
              calls.add(call);
          }
          return calls;
      }
  
      private void setQueryCondition(SearchBuilder search, Duration duration,
                                     List<String> serviceIds, String indexName) {
          final BoolQueryBuilder query =
              Query.bool()
                   .must(Query.range(ServiceRelationServerSideMetrics.TIME_BUCKET)
                              .gte(duration.getStartTimeBucket())
                              .lte(duration.getEndTimeBucket()));
  
          final BoolQueryBuilder serviceIdBoolQuery = Query.bool();
  
          query.must(serviceIdBoolQuery);
          if (IndexController.LogicIndicesRegister.isMergedTable(indexName)) {
              query.must(Query.term(IndexController.LogicIndicesRegister.METRIC_TABLE_NAME, indexName));
          }
          if (serviceIds.size() == 1) {
              serviceIdBoolQuery.should(
                  Query.term(
                      ServiceRelationServerSideMetrics.SOURCE_SERVICE_ID,
                      serviceIds.get(0)
                  ));
              serviceIdBoolQuery.should(
                  Query.term(
                      ServiceRelationServerSideMetrics.DEST_SERVICE_ID,
                      serviceIds.get(0)
                  ));
          } else {
              serviceIdBoolQuery.should(
                  Query.terms(
                      ServiceRelationServerSideMetrics.SOURCE_SERVICE_ID,
                      serviceIds
                  ));
              serviceIdBoolQuery.should(
                  Query.terms(
                      ServiceRelationServerSideMetrics.DEST_SERVICE_ID,
                      serviceIds
                  ));
          }
          search.query(query);
      }
  }
  