HIVE-29584: Invalidate RawStore after Direct SQL connection failures

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/DatabaseProduct.java

@@ -21,6 +21,9 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.SQLIntegrityConstraintViolationException;
+import java.sql.SQLNonTransientConnectionException;
+import java.sql.SQLRecoverableException;
+import java.sql.SQLTransientConnectionException;
 import java.sql.SQLTransactionRollbackException;
 import java.sql.Timestamp;
 import java.util.ArrayList;
@@ -217,6 +220,27 @@ public static boolean isRecoverableException(Throwable t) {
                  .allMatch(ex -> ExceptionUtils.indexOfType(t, ex) < 0);
   }
 
+  public static boolean isConnectionException(Throwable t) {
+    while (t != null) {
+      if (t instanceof SQLNonTransientConnectionException
+          || t instanceof SQLRecoverableException
+          || t instanceof SQLTransientConnectionException) {
+        return true;
+      }
+      if (t instanceof SQLException sqlException) {
+        for (SQLException current = sqlException; current != null;
+            current = current.getNextException()) {
+          String sqlState = current.getSQLState();
+          if (sqlState != null && sqlState.startsWith("08")) {
+            return true;
+          }
+        }
+      }
+      t = t.getCause();
+    }
+    return false;
+  }
+
   public final boolean isDERBY() {
     return dbType == DbType.DERBY;
   }

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/HMSHandlerContext.java

@@ -155,6 +155,21 @@ public static void setRawStore(RawStore rawStore) {
     context.get().rawStore = rawStore;
   }
 
+  /**
+   * Remove and shut down the RawStore associated with the current handler thread.
+   *
+   * @return true if a RawStore was present
+   */
+  public static boolean invalidateRawStore() {
+    RawStore rawStore = context.get().rawStore;
+    context.get().rawStore = null;
+    if (rawStore == null) {
+      return false;
+    }
+    rawStore.shutdown();
+    return true;
+  }
+
   public static void setTxnStore(TxnStore txnStore) {
     context.get().txnStore = txnStore;
   }

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/ObjectStore.java

@@ -377,9 +377,13 @@ public PersistenceManager getPersistenceManager() {
   @Override
   public void shutdown() {
     LOG.info("RawStore: {}, with PersistenceManager: {} will be shutdown", this, pm);
-    if (pm != null) {
-      pm.close();
-      pm = null;
+    PersistenceManager persistenceManager = pm;
+    pm = null;
+    currentTransaction = null;
+    openTrasactionCalls = 0;
+    transactionStatus = TXN_STATUS.NO_STATE;
+    if (persistenceManager != null) {
+      persistenceManager.close();
     }
   }
 
@@ -548,7 +552,9 @@ public void rollbackTransaction() {
       // remove all detached objects from the cache, since the transaction is
       // being rolled back they are no longer relevant, and this prevents them
       // from reattaching in future transactions
-      pm.evictAll();
+      if (pm != null) {
+        pm.evictAll();
+      }
     }
   }
 

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/directsql/MetaStoreDirectSql.java

@@ -1983,7 +1983,9 @@ public void prepareTxn() throws MetaException {
       assert pm.currentTransaction().isActive(); // must be inside tx together with queries
       executeNoResult(stmt);
     } catch (SQLException sqlEx) {
-      throw new MetaException("Error setting ansi quotes: " + sqlEx.getMessage());
+      MetaException metaException = new MetaException("Error setting ansi quotes: " + sqlEx.getMessage());
+      metaException.initCause(sqlEx);
+      throw metaException;
     }
   }
 

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/metastore/GetHelper.java

@@ -21,13 +21,15 @@
 import com.codahale.metrics.Counter;
 import com.google.common.annotations.VisibleForTesting;
 
+import javax.jdo.JDODataStoreException;
 import javax.jdo.JDOException;
 import javax.jdo.PersistenceManager;
 import java.util.List;
 
 import org.apache.hadoop.hive.common.TableName;
 import org.apache.hadoop.hive.metastore.DatabaseProduct;
 import org.apache.hadoop.hive.metastore.ExceptionHandler;
+import org.apache.hadoop.hive.metastore.HMSHandlerContext;
 import org.apache.hadoop.hive.metastore.directsql.MetaStoreDirectSql;
 import org.apache.hadoop.hive.metastore.RawStore;
 import org.apache.hadoop.hive.metastore.api.InvalidInputException;
@@ -59,6 +61,7 @@ public abstract class GetHelper<A, T> {
   protected final List<String> partitionFields;
   protected final A argument;
   private boolean success = false;
+  private boolean storeInvalidated = false;
   protected T results = null;
 
   public GetHelper(RawStoreBundle rsb, A args) throws MetaException {
@@ -148,6 +151,14 @@ private void start(boolean initTable) throws MetaException, NoSuchObjectExceptio
   }
 
   private void handleDirectSqlError(Exception ex, String savePoint) throws MetaException, NoSuchObjectException {
+    if (DatabaseProduct.isConnectionException(ex)) {
+      LOG.warn("Direct SQL failed because the metastore database connection is not usable", ex);
+      directSqlErrors.inc();
+      invalidateRawStore(ex);
+      throw ExceptionHandler.newMetaException(new JDODataStoreException(
+          "Direct SQL failed because the metastore database connection is not usable", ex));
+    }
+
     String message = null;
     try {
       message = generateShorterMessage(ex);
@@ -195,6 +206,18 @@ private void handleDirectSqlError(Exception ex, String savePoint) throws MetaExc
     doUseDirectSql = false;
   }
 
+  private void invalidateRawStore(Exception originalException) {
+    storeInvalidated = true;
+    try {
+      if (!HMSHandlerContext.invalidateRawStore()) {
+        baseStore.shutdown();
+      }
+    } catch (Exception cleanupException) {
+      originalException.addSuppressed(cleanupException);
+      LOG.warn("Failed to shut down RawStore after a Direct SQL connection error", cleanupException);
+    }
+  }
+
   public void setTransactionSavePoint(String savePoint) {
     if (savePoint != null) {
       ((JDOTransaction) pm.currentTransaction()).setSavepoint(savePoint);
@@ -252,7 +275,7 @@ private T commit() {
   }
 
   private void close() {
-    if (!success) {
+    if (!success && !storeInvalidated) {
       baseStore.rollbackTransaction();
     }
   }
@@ -278,4 +301,4 @@ public static Counter setDirectSqlErrors(Counter counter) {
     directSqlErrors = counter;
     return counter;
   }
-}
+}

standalone-metastore/metastore-server/src/test/java/org/apache/hadoop/hive/metastore/TestObjectStore.java

@@ -100,6 +100,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.jdo.JDODataStoreException;
 import javax.jdo.Query;
 import java.sql.Connection;
 import java.sql.DriverManager;
@@ -1424,6 +1425,47 @@ protected Database getJdoResult() throws MetaException,
     Assert.assertEquals(1, directSqlErrors.getCount());
   }
 
+  @Test
+  public void testDirectSqlConnectionErrorInvalidatesRawStore() throws Exception {
+    AtomicBoolean runJdo = new AtomicBoolean(false);
+    MetaException directSqlFailure = null;
+    objectStore.openTransaction();
+    HMSHandlerContext.setRawStore(objectStore);
+    try {
+      new GetHelper<DatabaseName, Database>(objectStore.createRawStoreBundle(),
+          new DatabaseName(DEFAULT_CATALOG_NAME, "foo")) {
+        @Override
+        protected Database getSqlResult() throws MetaException {
+          MetaException ex = new MetaException("Direct SQL connection failure");
+          ex.initCause(new SQLException("Communication link failure", "08S01"));
+          throw ex;
+        }
+
+        @Override
+        protected String describeResult() {
+          return "";
+        }
+
+        @Override
+        protected Database getJdoResult() {
+          runJdo.set(true);
+          return null;
+        }
+      }.run(false);
+      Assert.fail("Expected a connection-level Direct SQL failure");
+    } catch (MetaException ex) {
+      directSqlFailure = ex;
+      Assert.assertFalse(HMSHandlerContext.getRawStore().isPresent());
+    } finally {
+      // Handlers that own an outer transaction still execute this cleanup after GetHelper fails.
+      objectStore.rollbackTransaction();
+      HMSHandlerContext.clear();
+    }
+
+    Assert.assertTrue(directSqlFailure.getCause() instanceof JDODataStoreException);
+    Assert.assertFalse(runJdo.get());
+  }
+
   @Test(expected = MetaException.class)
   public void testLockDbTableThrowsExceptionWhenTableIsNotAllowedToLock() throws Exception {
     MetaStoreDirectSql metaStoreDirectSql = new MetaStoreDirectSql(objectStore.getPersistenceManager(), conf, null);
@@ -2195,4 +2237,3 @@ AutoCloseable directsql(boolean tryDirectSql) {
     };
   }
 }
-