content: /**
   * 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
   * <p>
   * http://www.apache.org/licenses/LICENSE-2.0
   * <p>
   * 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.atlas.tools;
  
  import org.apache.atlas.ApplicationProperties;
  import org.apache.atlas.AtlasClientV2;
  import org.apache.atlas.AtlasException;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
  import org.apache.atlas.repository.graphdb.AtlasVertex;
  import org.apache.atlas.repository.graphdb.janus.AtlasJanusGraphDatabase;
  import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
  import org.apache.atlas.utils.AuthenticationUtil;
  import org.apache.atlas.utils.SSLUtil;
  import org.apache.commons.cli.CommandLine;
  import org.apache.commons.cli.DefaultParser;
  import org.apache.commons.cli.Options;
  import org.apache.commons.cli.ParseException;
  import org.apache.commons.configuration.Configuration;
  import org.janusgraph.core.JanusGraph;
  import org.janusgraph.core.schema.JanusGraphIndex;
  import org.janusgraph.core.schema.SchemaAction;
  import org.janusgraph.core.schema.SchemaStatus;
  import org.janusgraph.diskstorage.BackendTransaction;
  import org.janusgraph.diskstorage.indexing.IndexEntry;
  import org.janusgraph.graphdb.database.IndexSerializer;
  import org.janusgraph.graphdb.database.StandardJanusGraph;
  import org.janusgraph.graphdb.database.management.ManagementSystem;
  import org.janusgraph.graphdb.transaction.StandardJanusGraphTx;
  import org.janusgraph.graphdb.types.MixedIndexType;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.function.Consumer;
  
  public class RepairIndex {
      private static final Logger LOG = LoggerFactory.getLogger(RepairIndex.class);
  
      private static final int EXIT_CODE_SUCCESS                      = 0;
      private static final int EXIT_CODE_FAILED                       = 1;
      private static final int MAX_TRIES_ON_FAILURE                   = 3;
      private static final String INDEX_NAME_VERTEX_INDEX             = "vertex_index";
      private static final String INDEX_NAME_FULLTEXT_INDEX           = "fulltext_index";
      private static final String INDEX_NAME_EDGE_INDEX               = "edge_index";
      private static final String DEFAULT_ATLAS_URL                   = "http://localhost:21000/";
      private static final String APPLICATION_PROPERTY_ATLAS_ENDPOINT = "atlas.rest.address";
  
      private static JanusGraph    graph;
      private static AtlasClientV2 atlasClientV2;
      private static boolean       isSelectiveRestore;
  
      public static void main(String[] args) {
          int exitCode = EXIT_CODE_FAILED;
          LOG.info("Started index repair");
  
          try {
              CommandLine cmd  = getCommandLine(args);
              String      guid = cmd.getOptionValue("g");
  
              if (guid != null && !guid.isEmpty()) {
                  isSelectiveRestore = true;
                  String uid = cmd.getOptionValue("u");
                  String pwd = cmd.getOptionValue("p");
                  setupAtlasClient(uid, pwd);
              }
  
              process(guid);
  
              LOG.info("Completed index repair!");
              exitCode = EXIT_CODE_SUCCESS;
          } catch (Exception e) {
              LOG.error("Failed!", e);
              display("Failed: " + e.getMessage());
          }
  
          System.exit(exitCode);
      }
  
      private static void process(String guid) throws Exception {
          RepairIndex repairIndex = new RepairIndex();
  
          SSLUtil sslUtil = new SSLUtil();
          sslUtil.setSSLContext();
  
          setupGraph();
  
          if (isSelectiveRestore) {
              repairIndex.restoreSelective(guid);
          } else {
              repairIndex.restoreAll();
          }
  
          displayCrlf("Repair Index: Done!");
      }
  
      private static CommandLine getCommandLine(String[] args) throws ParseException {
          Options options = new Options();
          options.addOption("g", "guid", true, "guid for which update index should be executed.");
          options.addOption("u", "user", true, "User name.");
          options.addOption("p", "password", true, "Password name.");
  
          return new DefaultParser().parse(options, args);
      }
  
      private static void setupGraph() {
          display("Initializing graph: ");
          graph = AtlasJanusGraphDatabase.getGraphInstance();
          displayCrlf("Graph Initialized!");
      }
  
      private static String[] getIndexes() {
          return new String[] {INDEX_NAME_VERTEX_INDEX, INDEX_NAME_EDGE_INDEX, INDEX_NAME_FULLTEXT_INDEX};
      }
  
      private static void setupAtlasClient(String uid, String pwd) throws AtlasException {
          String[] atlasEndpoint = getAtlasRESTUrl();
          if (atlasEndpoint == null || atlasEndpoint.length == 0) {
              atlasEndpoint = new String[] {DEFAULT_ATLAS_URL};
          }
          atlasClientV2 = getAtlasClientV2(atlasEndpoint, new String[] {uid, pwd});
      }
  
      private void restoreAll() throws Exception {
          for (String indexName : getIndexes()) {
              displayCrlf("Restoring: " + indexName);
  
              long             startTime = System.currentTimeMillis();
              ManagementSystem mgmt      = null;
  
              try {
                  mgmt = (ManagementSystem) graph.openManagement();
  
                  JanusGraphIndex index = mgmt.getGraphIndex(indexName);
  
                  mgmt.updateIndex(index, SchemaAction.REINDEX).get();
              } finally {
                  if (mgmt != null) {
                      mgmt.commit();
                  }
              }
  
              ManagementSystem.awaitGraphIndexStatus(graph, indexName).status(SchemaStatus.ENABLED).call();
  
              display(": Time taken: " + (System.currentTimeMillis() - startTime) + " ms");
              displayCrlf(": Done!");
          }
      }
  
      private void restoreSelective(String guid) throws Exception {
          Set<String> referencedGUIDs = new HashSet<>(getEntityAndReferenceGuids(guid));
          displayCrlf("processing referencedGuids => " + referencedGUIDs);
  
          StandardJanusGraph janusGraph      = (StandardJanusGraph) graph;
          IndexSerializer    indexSerializer = janusGraph.getIndexSerializer();
  
          for (String indexName : getIndexes()) {
              displayCrlf("Restoring: " + indexName);
              long startTime = System.currentTimeMillis();
              reindexVertex(indexName, indexSerializer, referencedGUIDs);
  
              display(": Time taken: " + (System.currentTimeMillis() - startTime) + " ms");
              displayCrlf(": Done!");
          }
      }
  
      private static void reindexVertex(String indexName, IndexSerializer indexSerializer, Set<String> entityGUIDs) throws Exception {
          ManagementSystem mgmt = null;
  
          try {
              mgmt = (ManagementSystem) graph.openManagement();
  
              StandardJanusGraphTx tx        = mgmt.getWrappedTx();
              BackendTransaction   mutator   = tx.getTxHandle();
              JanusGraphIndex      index     = mgmt.getGraphIndex(indexName);
              MixedIndexType       indexType = (MixedIndexType) mgmt.getSchemaVertex(index).asIndexType();
  
              Map<String, Map<String, List<IndexEntry>>> documentsPerStore = new java.util.HashMap<>();
  
              for (String entityGuid : entityGUIDs) {
                  for (int attemptCount = 1; attemptCount <= MAX_TRIES_ON_FAILURE; attemptCount++) {
                      AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(entityGuid);
                      try {
                          indexSerializer.reindexElement(vertex.getWrappedElement(), indexType, documentsPerStore);
                          break;
                      } catch (Exception e) {
                          displayCrlf("Exception: " + e.getMessage());
                          displayCrlf("Pausing before retry..");
                          Thread.sleep(2000 * attemptCount);
                      }
                  }
              }
              mutator.getIndexTransaction(indexType.getBackingIndexName()).restore(documentsPerStore);
          } finally {
              if (mgmt != null) {
                  mgmt.commit();
              }
          }
      }
  
      private static Set<String> getEntityAndReferenceGuids(String guid) throws Exception {
          Set<String> set = new HashSet<>();
          set.add(guid);
          AtlasEntityWithExtInfo   entity = atlasClientV2.getEntityByGuid(guid);
          Map<String, AtlasEntity> map    = entity.getReferredEntities();
          if (map == null || map.isEmpty()) {
              return set;
          }
          set.addAll(map.keySet());
          return set;
      }
  
      private static void display(String... formatMessage) {
          displayFn(System.out::print, formatMessage);
      }
  
      private static void displayCrlf(String... formatMessage) {
          displayFn(System.out::println, formatMessage);
      }
  
      private static void displayFn(Consumer<String> fn, String... formatMessage) {
          if (formatMessage.length == 1) {
              fn.accept(formatMessage[0]);
          } else {
              fn.accept(String.format(formatMessage[0], formatMessage[1]));
          }
      }
  
      private static String[] getAtlasRESTUrl() {
          Configuration atlasConf;
          try {
              atlasConf = ApplicationProperties.get();
              return atlasConf.getStringArray(APPLICATION_PROPERTY_ATLAS_ENDPOINT);
          } catch (AtlasException e) {
              return new String[] {DEFAULT_ATLAS_URL};
          }
      }
  
      private static AtlasClientV2 getAtlasClientV2(String[] atlasEndpoint, String[] uidPwdFromCommandLine) throws AtlasException {
          AtlasClientV2 atlasClientV2;
          if (!AuthenticationUtil.isKerberosAuthenticationEnabled()) {
              String[] uidPwd = (uidPwdFromCommandLine[0] == null || uidPwdFromCommandLine[1] == null) ? AuthenticationUtil.getBasicAuthenticationInput() : uidPwdFromCommandLine;
  
              atlasClientV2 = new AtlasClientV2(atlasEndpoint, uidPwd);
          } else {
              atlasClientV2 = new AtlasClientV2(atlasEndpoint);
          }
          return atlasClientV2;
      }
  }
  