From 6b23e02d2bf3f67ed9619453e65a2299a20a1c76 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Xavier=20L=C3=A9aut=C3=A9?= <xavier@metamarkets.com>
Date: Fri, 22 May 2015 09:35:03 -0400
Subject: [PATCH] retry transient exceptions for PostgreSQL, fixes #1382

---
 .../postgresql-metadata-storage/pom.xml       |  6 +++
 .../postgresql/PostgreSQLConnector.java       | 15 ++++++
 .../postgresql/PostgreSQLConnectorTest.java   | 49 +++++++++++++++++++
 3 files changed, 70 insertions(+)
 create mode 100644 extensions/postgresql-metadata-storage/src/test/java/io/druid/metadata/storage/postgresql/PostgreSQLConnectorTest.java

diff --git a/extensions/postgresql-metadata-storage/pom.xml b/extensions/postgresql-metadata-storage/pom.xml
index 14529ee1b0d..ed3aa9025cc 100644
--- a/extensions/postgresql-metadata-storage/pom.xml
+++ b/extensions/postgresql-metadata-storage/pom.xml
@@ -56,6 +56,12 @@
             <groupId>org.jdbi</groupId>
             <artifactId>jdbi</artifactId>
         </dependency>
+
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 </project>
diff --git a/extensions/postgresql-metadata-storage/src/main/java/io/druid/metadata/storage/postgresql/PostgreSQLConnector.java b/extensions/postgresql-metadata-storage/src/main/java/io/druid/metadata/storage/postgresql/PostgreSQLConnector.java
index 1ffd0248412..6504e938b13 100644
--- a/extensions/postgresql-metadata-storage/src/main/java/io/druid/metadata/storage/postgresql/PostgreSQLConnector.java
+++ b/extensions/postgresql-metadata-storage/src/main/java/io/druid/metadata/storage/postgresql/PostgreSQLConnector.java
@@ -29,6 +29,8 @@ import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.skife.jdbi.v2.util.StringMapper;
 
+import java.sql.SQLException;
+
 public class PostgreSQLConnector extends SQLMetadataConnector
 {
   private static final Logger log = new Logger(PostgreSQLConnector.class);
@@ -112,4 +114,17 @@ public class PostgreSQLConnector extends SQLMetadataConnector
 
   @Override
   public DBI getDBI() { return dbi; }
+
+  @Override
+  protected boolean isTransientException(Throwable e)
+  {
+    if(e instanceof SQLException) {
+      final String sqlState = ((SQLException) e).getSQLState();
+      // limited to errors that are likely to be resolved within a few retries
+      // retry on connection errors and insufficient resources
+      // see http://www.postgresql.org/docs/current/static/errcodes-appendix.html for details
+      return sqlState != null && (sqlState.startsWith("08") || sqlState.startsWith("53"));
+    }
+    return false;
+  }
 }
diff --git a/extensions/postgresql-metadata-storage/src/test/java/io/druid/metadata/storage/postgresql/PostgreSQLConnectorTest.java b/extensions/postgresql-metadata-storage/src/test/java/io/druid/metadata/storage/postgresql/PostgreSQLConnectorTest.java
new file mode 100644
index 00000000000..43f1e34e9f6
--- /dev/null
+++ b/extensions/postgresql-metadata-storage/src/test/java/io/druid/metadata/storage/postgresql/PostgreSQLConnectorTest.java
@@ -0,0 +1,49 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  Metamarkets 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 io.druid.metadata.storage.postgresql;
+
+import com.google.common.base.Suppliers;
+import io.druid.metadata.MetadataStorageConnectorConfig;
+import io.druid.metadata.MetadataStorageTablesConfig;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.sql.SQLException;
+
+public class PostgreSQLConnectorTest
+{
+
+  @Test
+  public void testIsTransientException() throws Exception
+  {
+    PostgreSQLConnector connector = new PostgreSQLConnector(
+        Suppliers.ofInstance(new MetadataStorageConnectorConfig()),
+        Suppliers.ofInstance(new MetadataStorageTablesConfig(null, null, null, null, null, null, null, null))
+    );
+
+    Assert.assertTrue(connector.isTransientException(new SQLException("bummer, connection problem", "08DIE")));
+    Assert.assertTrue(connector.isTransientException(new SQLException("bummer, too many things going on", "53RES")));
+    Assert.assertFalse(connector.isTransientException(new SQLException("oh god, no!", "58000")));
+    Assert.assertFalse(connector.isTransientException(new SQLException("help!")));
+    Assert.assertFalse(connector.isTransientException(new SQLException()));
+    Assert.assertFalse(connector.isTransientException(new Exception("I'm not happy")));
+    Assert.assertFalse(connector.isTransientException(new Throwable("I give up")));
+  }
+}