fix existDatabase(), clearBackend() and rollback() for postgresql

fixed: #530 Change-Id: I5061eb95abf0e0ef3a6a83a74ced291647d9308d
apache · May 28, 2019 · f88d047 · f88d047
1 parent 0061d4f
commit f88d047
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/config/CoreOptions.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/config/CoreOptions.java
@@ -125,7 +125,7 @@ public static synchronized CoreOptions instance() {
 
     public static final ConfigOption<Boolean> VERTEX_CHECK_CUSTOMIZED_ID_EXIST =
             new ConfigOption<>(
-                    "vertex.check_customzied_id_exist",
+                    "vertex.check_customized_id_exist",
                     "Whether to check the vertices exist for those using " +
                     "customized id strategy",
                     disallowEmpty(),

diff --git a/hugegraph-mysql/src/main/java/com/baidu/hugegraph/backend/store/mysql/MysqlSessions.java b/hugegraph-mysql/src/main/java/com/baidu/hugegraph/backend/store/mysql/MysqlSessions.java
@@ -168,8 +168,7 @@ protected String buildCreateDatabase(String database) {
     public void dropDatabase() {
         LOG.debug("Drop database: {}", this.database);
 
-        String sql = String.format("DROP DATABASE IF EXISTS %s;",
-                                   this.database);
+        String sql = this.buildDropDatabase(this.database);
         try (Connection conn = this.openWithoutDB(DROP_DB_TIMEOUT)) {
             conn.createStatement().execute(sql);
         } catch (SQLException e) {
@@ -182,6 +181,10 @@ public void dropDatabase() {
         }
     }
 
+    protected String buildDropDatabase(String database) {
+        return String.format("DROP DATABASE IF EXISTS %s;", database);
+    }
+
     public boolean existsDatabase() {
         try (Connection conn = this.openWithoutDB(0);
              ResultSet result = conn.getMetaData().getCatalogs()) {
@@ -312,12 +315,9 @@ public Integer commit() {
                 }
                 this.conn.commit();
                 this.clear();
+                this.end();
             } catch (SQLException e) {
                 throw new BackendException("Failed to commit", e);
-            } finally {
-                try {
-                    this.end();
-                } catch (SQLException ignored) {}
             }
             return updated;
         }

diff --git a/...gresql/src/main/java/com/baidu/hugegraph/backend/store/postgresql/PostgresqlSessions.java b/...gresql/src/main/java/com/baidu/hugegraph/backend/store/postgresql/PostgresqlSessions.java
@@ -20,6 +20,7 @@
 package com.baidu.hugegraph.backend.store.postgresql;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 
 import org.apache.http.client.utils.URIBuilder;
@@ -45,6 +46,20 @@ public PostgresqlSessions(HugeConfig config, String database, String store) {
         super(config, database, store);
     }
 
+    @Override
+    public boolean existsDatabase() {
+        String statement = String.format(
+                           "SELECT datname FROM pg_catalog.pg_database " +
+                           "WHERE datname = '%s';", this.database());
+        try (Connection conn = this.openWithoutDB(0)) {
+            ResultSet result = conn.createStatement().executeQuery(statement);
+            return result.next();
+        } catch (Exception e) {
+            throw new BackendException("Failed to obtain MySQL metadata, " +
+                                       "please ensure it is ok", e);
+        }
+    }
+
     @Override
     public void createDatabase() {
         // Create database with non-database-session
@@ -76,6 +91,17 @@ protected String buildCreateDatabase(String database) {
         return String.format(POSTGRESQL_DB_CREATE, database);
     }
 
+    @Override
+    protected String buildDropDatabase(String database) {
+        return String.format(
+               "REVOKE CONNECT ON DATABASE %s FROM public;" +
+               "SELECT pg_terminate_backend(pg_stat_activity.pid) " +
+               "  FROM pg_stat_activity " +
+               "  WHERE pg_stat_activity.datname = '%s';" +
+               "DROP DATABASE IF EXISTS %s;",
+               database, database, database);
+    }
+
     @Override
     protected URIBuilder newConnectionURIBuilder() {
         // Suppress error log when database does not exist