Optimize load database method in JdbcDatabaseOperations #6629

Author: sunxiaojian

catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/DatabaseOperation.java

@@ -58,6 +58,14 @@ void create(String databaseName, String comment, Map<String, String> properties)
   /** @return The list name of databases. */
   List<String> listDatabases();
 
+  /**
+   * Checks if the specified database exists.
+   *
+   * @param databaseName The name of the database to check.
+   * @return true if the database exists; false otherwise.
+   */
+  boolean exist(String databaseName);
+
   /**
    * @param databaseName The name of the database to check.
    * @return information object of the JDBC database.

catalogs/catalog-jdbc-common/src/main/java/org/apache/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java

@@ -107,6 +107,28 @@ public List<String> listDatabases() {
     }
   }
 
+  @Override
+  public boolean exist(String databaseName) {
+    try (final Connection connection = this.dataSource.getConnection()) {
+      String query = generateDatabaseExistSql(databaseName);
+      try (Statement statement = connection.createStatement()) {
+        try (ResultSet resultSet = statement.executeQuery(query)) {
+          if (resultSet.next()) {
+            return true;
+          }
+        }
+      }
+    } catch (SQLException sqlException) {
+      throw this.exceptionMapper.toGravitinoException(sqlException);
+    }
+    return false;
+  }
+
+  protected String generateDatabaseExistSql(String databaseName) {
+    return String.format(
+        "SELECT * FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = '%s'", databaseName);
+  }
+
   protected void dropDatabase(String databaseName, boolean cascade) {
     try (final Connection connection = getConnection()) {
       JdbcConnectorUtils.executeUpdate(connection, generateDropDatabaseSql(databaseName, cascade));
@@ -176,16 +198,11 @@ protected String generateDropDatabaseSql(String databaseName, boolean cascade) {
    */
   @Override
   public JdbcSchema load(String databaseName) throws NoSuchSchemaException {
-    List<String> allDatabases = listDatabases();
-    String dbName =
-        allDatabases.stream()
-            .filter(db -> db.equals(databaseName))
-            .findFirst()
-            .orElseThrow(
-                () -> new NoSuchSchemaException("Database %s could not be found", databaseName));
-
+    if (!exist(databaseName)) {
+      throw new NoSuchSchemaException("Database %s could not be found", databaseName);
+    }
     return JdbcSchema.builder()
-        .withName(dbName)
+        .withName(databaseName)
         .withProperties(ImmutableMap.of())
         .withAuditInfo(AuditInfo.EMPTY)
         .build();

catalogs/catalog-jdbc-common/src/test/java/org/apache/gravitino/catalog/jdbc/operation/SqliteDatabaseOperations.java

@@ -67,6 +67,7 @@ public void create(String databaseName, String comment, Map<String, String> prop
     Preconditions.checkArgument(exist(databaseName), "Database %s does not exist", databaseName);
   }
 
+  @Override
   public boolean exist(String databaseName) {
     return new File(dbPath + "/" + databaseName).exists();
   }

catalogs/catalog-jdbc-postgresql/src/main/java/org/apache/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java

@@ -155,6 +155,12 @@ protected Connection getConnection() throws SQLException {
     return connection;
   }
 
+  @Override
+  protected String generateDatabaseExistSql(String databaseName) {
+    return String.format(
+        "SELECT n.datname FROM pg_catalog.pg_database n where n.datname='%s'", databaseName);
+  }
+
   @Override
   protected boolean supportSchemaComment() {
     return true;