[WIP] HIVE-29668: Add -rebuildIndexes utility to reconstruct backend Metastore indexes

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/AbstractIndexRebuilder.java

@@ -0,0 +1,81 @@
+/*
+ * 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.hadoop.hive.metastore.tools.schematool;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.stream.Collectors;
+
+import org.apache.hadoop.hive.metastore.HiveMetaException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Shared duplicate-check and DDL execution logic for {@link IndexRebuilder}. */
+public abstract class AbstractIndexRebuilder implements IndexRebuilder {
+
+  private static final Logger LOG = LoggerFactory.getLogger(AbstractIndexRebuilder.class);
+
+  protected final Connection conn;
+  protected final boolean needsQuotedIdentifier;
+  protected final String quoteCharacter;
+
+  protected AbstractIndexRebuilder(Connection conn, boolean needsQuotedIdentifier,
+      String quoteCharacter) {
+    this.conn = conn;
+    this.needsQuotedIdentifier = needsQuotedIdentifier;
+    this.quoteCharacter = quoteCharacter;
+  }
+
+  @Override
+  public long findDuplicates(IndexInfo index) throws HiveMetaException {
+    if (!index.unique()) {
+      return 0;
+    }
+    String quotedCols = index.columns().stream()
+        .map(c -> "<qa>" + c + "<qa>")
+        .collect(Collectors.joining(", "));
+    String sql = MetastoreSchemaTool.quote(
+        "SELECT COUNT(*) FROM (SELECT " + quotedCols
+            + " FROM <qa>" + index.tableName() + "<qa>"
+            + " GROUP BY " + quotedCols
+            + " HAVING COUNT(*) > 1) duplicates",
+        needsQuotedIdentifier, quoteCharacter);
+    try (PreparedStatement ps = conn.prepareStatement(sql);
+        ResultSet rs = ps.executeQuery()) {
+      return rs.next() ? rs.getLong(1) : 0;
+    } catch (SQLException e) {
+      throw new HiveMetaException(
+          "Failed to check for duplicate rows for index \"" + index.indexName() + "\"", e);
+    }
+  }
+
+  /** Executes one or more DDL statements, logging each before execution. */
+  protected void executeRebuild(IndexInfo index, String... ddls) throws HiveMetaException {
+    try (Statement stmt = conn.createStatement()) {
+      for (String ddl : ddls) {
+        LOG.info("Executing: {}", ddl);
+        stmt.execute(ddl);
+      }
+    } catch (SQLException e) {
+      throw new HiveMetaException("Failed to rebuild index \"" + index.indexName() + "\"", e);
+    }
+  }
+}

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/HiveSchemaHelper.java

@@ -42,7 +42,7 @@ public class HiveSchemaHelper {
   public static final String DB_HIVE = "hive";
   public static final String DB_MSSQL = "mssql";
   public static final String DB_MYSQL = "mysql";
-  public static final String DB_POSTGRACE = "postgres";
+  public static final String DB_POSTGRES = "postgres";
   public static final String DB_ORACLE = "oracle";
   public static final String EMBEDDED_HS2_URL =
       "jdbc:hive2://?hive.conf.restricted.list=;hive.security.authorization.sqlstd.confwhitelist=.*;"
@@ -581,7 +581,7 @@ public static NestedScriptParser getDbCommandParser(String dbName,
       return new MSSQLCommandParser(dbOpts, msUsername, msPassword, conf, usingSqlLine);
     } else if (dbName.equalsIgnoreCase(DB_MYSQL)) {
       return new MySqlCommandParser(dbOpts, msUsername, msPassword, conf, usingSqlLine);
-    } else if (dbName.equalsIgnoreCase(DB_POSTGRACE)) {
+    } else if (dbName.equalsIgnoreCase(DB_POSTGRES)) {
       return new PostgresCommandParser(dbOpts, msUsername, msPassword, conf, usingSqlLine);
     } else if (dbName.equalsIgnoreCase(DB_ORACLE)) {
       return new OracleCommandParser(dbOpts, msUsername, msPassword, conf, usingSqlLine);

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/IndexInfo.java

@@ -0,0 +1,42 @@
+/*
+ * 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.hadoop.hive.metastore.tools.schematool;
+
+import java.util.List;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Record for a DB index, used by {@link IndexRebuilder}.
+ */
+public record IndexInfo(
+    String indexName,
+    String tableName,
+    boolean unique,
+    boolean constraintBacked,
+    List<String> columns) {
+
+  public IndexInfo {
+    columns = List.copyOf(columns);
+  }
+
+  @Override
+  public @NotNull String toString() {
+    return (constraintBacked ? "CONSTRAINT" : "INDEX")
+        + " \"" + indexName + "\" ON \"" + tableName + "\"";
+  }
+}

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/IndexRebuilder.java

@@ -0,0 +1,43 @@
+/*
+ * 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.hadoop.hive.metastore.tools.schematool;
+
+import java.util.List;
+
+import org.apache.hadoop.hive.metastore.HiveMetaException;
+
+/**
+ * Interface for rebuilding indexes in the HMS backend database.
+ * Implementations should extend {@link AbstractIndexRebuilder}.
+ */
+public interface IndexRebuilder {
+
+  /** Returns all B-tree indexes in the current HMS schema. */
+  List<IndexInfo> loadIndexes() throws HiveMetaException;
+
+  /**
+   * Returns duplicate key-group count for the index. Rebuild must be blocked when this is
+   * greater than zero. Always returns zero for non-unique indexes.
+   */
+  long findDuplicates(IndexInfo index) throws HiveMetaException;
+
+  void rebuildIndex(IndexInfo index) throws HiveMetaException;
+
+  /** Returns a description of the DDL that would be executed to rebuild the index. */
+  String describeRebuildDDL(IndexInfo index);
+}

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/IndexRebuilderFactory.java

@@ -0,0 +1,48 @@
+/*
+ * 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.hadoop.hive.metastore.tools.schematool;
+
+import java.sql.Connection;
+
+import org.apache.hadoop.hive.metastore.HiveMetaException;
+
+/**
+ * Factory for creating {@link IndexRebuilder} instances for the given database type.
+ * Add support for a new backend by adding a {@code case} branch.
+ */
+public final class IndexRebuilderFactory {
+
+  private IndexRebuilderFactory() {
+  }
+
+  public static IndexRebuilder create(String dbType, Connection conn,
+      MetastoreSchemaTool schemaTool) throws HiveMetaException {
+    return switch (dbType.toLowerCase()) {
+      case HiveSchemaHelper.DB_POSTGRES ->
+          new PostgresIndexRebuilder(conn, schemaTool.needsQuotedIdentifier, schemaTool.quoteCharacter);
+      case HiveSchemaHelper.DB_MYSQL ->
+          new MySQLIndexRebuilder(conn, schemaTool.needsQuotedIdentifier, schemaTool.quoteCharacter);
+      case HiveSchemaHelper.DB_ORACLE ->
+          new OracleIndexRebuilder(conn, schemaTool.needsQuotedIdentifier, schemaTool.quoteCharacter);
+      case HiveSchemaHelper.DB_MSSQL ->
+          new MSSQLIndexRebuilder(conn, schemaTool.needsQuotedIdentifier, schemaTool.quoteCharacter);
+      default -> throw new HiveMetaException(
+          "-rebuildIndexes is not supported for -dbType " + dbType + ".");
+    };
+  }
+}

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/MSSQLIndexRebuilder.java

@@ -0,0 +1,120 @@
+/*
+ * 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.hadoop.hive.metastore.tools.schematool;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.hadoop.hive.metastore.HiveMetaException;
+
+/**
+ * SQL Server implementation of {@link IndexRebuilder}.
+ *
+ * <p>Uses {@code sys.indexes} / {@code sys.index_columns} / {@code sys.columns} and rebuilds
+ * with {@code ALTER INDEX name ON table REBUILD}.
+ */
+class MSSQLIndexRebuilder extends AbstractIndexRebuilder {
+
+  // i.type: 1 = clustered, 2 = nonclustered.
+  // i.name IS NOT NULL excludes heap pseudo-entries (type 0).
+  private static final String QUERY_INDEXES = """
+      SELECT i.name  AS index_name,
+             t.name  AS table_name,
+             i.is_unique
+        FROM sys.indexes i
+        JOIN sys.tables t ON t.object_id = i.object_id
+       WHERE t.is_ms_shipped = 0
+         AND i.type IN (1, 2)
+         AND i.name IS NOT NULL
+       ORDER BY t.name, i.name""";
+
+  // INCLUDE columns are not key columns; exclude them.
+  private static final String QUERY_INDEX_COLUMNS = """
+      SELECT c.name AS column_name
+        FROM sys.indexes i
+        JOIN sys.tables t          ON t.object_id = i.object_id
+        JOIN sys.index_columns ic  ON ic.object_id = i.object_id
+                                  AND ic.index_id  = i.index_id
+        JOIN sys.columns c         ON c.object_id  = ic.object_id
+                                  AND c.column_id  = ic.column_id
+       WHERE t.name = ? AND i.name = ?
+         AND ic.is_included_column = 0
+       ORDER BY ic.key_ordinal""";
+
+  MSSQLIndexRebuilder(Connection conn, boolean needsQuotedIdentifier, String quoteCharacter) {
+    super(conn, needsQuotedIdentifier, quoteCharacter);
+  }
+
+  @Override
+  public List<IndexInfo> loadIndexes() throws HiveMetaException {
+    try (Statement stmt = conn.createStatement();
+        ResultSet rs = stmt.executeQuery(QUERY_INDEXES)) {
+      List<IndexInfo> indexes = new ArrayList<>();
+      while (rs.next()) {
+        String indexName = rs.getString("index_name");
+        String tableName = rs.getString("table_name");
+        boolean isUnique = rs.getBoolean("is_unique");
+        List<String> columns = loadIndexColumns(tableName, indexName);
+        // constraintBacked is irrelevant for MSSQL: ALTER INDEX REBUILD works identically
+        // for all index types
+        indexes.add(new IndexInfo(indexName, tableName, isUnique, false, columns));
+      }
+      return indexes;
+    } catch (SQLException e) {
+      throw new HiveMetaException("Failed to load indexes from SQL Server catalog", e);
+    }
+  }
+
+  @Override
+  public void rebuildIndex(IndexInfo index) throws HiveMetaException {
+    executeRebuild(index, buildRebuildDdl(index));
+  }
+
+  @Override
+  public String describeRebuildDDL(IndexInfo index) {
+    return buildRebuildDdl(index) + ";";
+  }
+
+  private String buildRebuildDdl(IndexInfo index) {
+    // MSSQL requires the table name because index names are unique per-table, not per-schema.
+    return MetastoreSchemaTool.quote(
+        "ALTER INDEX <qa>" + index.indexName() + "<qa>"
+            + " ON <qa>" + index.tableName() + "<qa> REBUILD",
+        needsQuotedIdentifier, quoteCharacter);
+  }
+
+  private List<String> loadIndexColumns(String tableName, String indexName) throws SQLException {
+    List<String> columns = new ArrayList<>();
+    try (PreparedStatement ps = conn.prepareStatement(QUERY_INDEX_COLUMNS)) {
+      ps.setString(1, tableName);
+      ps.setString(2, indexName);
+      try (ResultSet rs = ps.executeQuery()) {
+        while (rs.next()) {
+          columns.add(rs.getString("column_name"));
+        }
+      }
+    }
+    return columns;
+  }
+}
+

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/tools/schematool/MetastoreSchemaTool.java

@@ -477,6 +477,8 @@ public int run(String metastoreHome, String[] args, OptionGroup additionalOption
         task = new SchemaToolTaskDrop();
       } else if (cmdLine.hasOption("createLogsTable")) {
         task = new SchemaToolTaskCreateLogsTable();
+      } else if (cmdLine.hasOption("rebuildIndexes")) {
+        task = new SchemaToolTaskRebuildIndexes();
       } else {
         throw new HiveMetaException("No task defined!");
       }