Fix to avoid deleting non ScalarDB tables when dropping namespaces

Author: KodaiD

core/src/main/java/com/scalar/db/storage/jdbc/JdbcAdmin.java

@@ -481,6 +481,11 @@ private void deleteMetadataSchema(Connection connection) throws SQLException {
   @Override
   public void dropNamespace(String namespace) throws ExecutionException {
     try (Connection connection = dataSource.getConnection()) {
+      Set<String> remainingTables = getNamespaceTableNamesInternal(connection, namespace);
+      if (!remainingTables.isEmpty()) {
+        throw new IllegalStateException(
+            CoreError.NAMESPACE_NOT_EMPTY.buildMessage(namespace, remainingTables));
+      }
       execute(connection, rdbEngine.dropNamespaceSql(namespace));
       deleteFromNamespacesTable(connection, namespace);
       deleteNamespacesTableAndMetadataSchemaIfEmpty(connection);
@@ -703,6 +708,20 @@ public Set<String> getNamespaceTableNames(String namespace) throws ExecutionExce
     }
   }
 
+  private Set<String> getNamespaceTableNamesInternal(Connection connection, String namespace)
+      throws SQLException {
+    String sql = rdbEngine.getTableNamesInNamespaceSql(namespace);
+    Set<String> tableNames = new HashSet<>();
+    try (PreparedStatement statement = connection.prepareStatement(sql)) {
+      try (ResultSet resultSet = statement.executeQuery()) {
+        while (resultSet.next()) {
+          tableNames.add(resultSet.getString(1));
+        }
+      }
+    }
+    return tableNames;
+  }
+
   @Override
   public boolean namespaceExists(String namespace) throws ExecutionException {
     String selectQuery =

core/src/main/java/com/scalar/db/storage/jdbc/RdbEngineDb2.java

@@ -583,4 +583,10 @@ public void throwIfCrossPartitionScanOrderingOnBlobColumnNotSupported(
               .buildMessage(orderingOnBlobColumn.get()));
     }
   }
+
+  @Override
+  public String getTableNamesInNamespaceSql(String namespace) {
+    return String.format(
+        "SELECT TABNAME FROM SYSCAT.TABLES WHERE TABSCHEMA = '%s' AND TYPE = 'T'", namespace);
+  }
 }

core/src/main/java/com/scalar/db/storage/jdbc/RdbEngineMysql.java

@@ -504,4 +504,11 @@ public void setConnectionToReadOnly(Connection connection, boolean readOnly) thr
     // Observed performance degradation when using read-only connections in MySQL. So we do not
     // set the read-only mode for MySQL connections.
   }
+
+  @Override
+  public String getTableNamesInNamespaceSql(String namespace) {
+    return "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = '"
+        + namespace
+        + "'";
+  }
 }

core/src/main/java/com/scalar/db/storage/jdbc/RdbEngineOracle.java

@@ -451,4 +451,9 @@ public void throwIfAlterColumnTypeNotSupported(DataType from, DataType to) {
               from.toString(), to.toString()));
     }
   }
+
+  @Override
+  public String getTableNamesInNamespaceSql(String namespace) {
+    return String.format("SELECT TABLE_NAME FROM ALL_TABLES WHERE OWNER = '%s'", namespace);
+  }
 }

core/src/main/java/com/scalar/db/storage/jdbc/RdbEnginePostgresql.java

@@ -395,4 +395,11 @@ public String tryAddIfNotExistsToCreateIndexSql(String createIndexSql) {
       getTimeTypeStrategy() {
     return timeTypeEngine;
   }
+
+  @Override
+  public String getTableNamesInNamespaceSql(String namespace) {
+    return "SELECT table_name FROM information_schema.tables WHERE table_schema = '"
+        + namespace
+        + "'";
+  }
 }

core/src/main/java/com/scalar/db/storage/jdbc/RdbEngineSqlServer.java

@@ -431,4 +431,11 @@ public Map<String, String> getConnectionProperties(JdbcConfig config) {
       getTimeTypeStrategy() {
     return timeTypeEngine;
   }
+
+  @Override
+  public String getTableNamesInNamespaceSql(String namespace) {
+    return "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = '"
+        + namespace
+        + "'";
+  }
 }

core/src/main/java/com/scalar/db/storage/jdbc/RdbEngineSqlite.java

@@ -367,4 +367,12 @@ public void throwIfAlterColumnTypeNotSupported(DataType from, DataType to) {
   public void setConnectionToReadOnly(Connection connection, boolean readOnly) {
     // Do nothing. SQLite does not support read-only mode.
   }
+
+  @Override
+  public String getTableNamesInNamespaceSql(String namespace) {
+    return "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE '"
+        + namespace
+        + NAMESPACE_SEPARATOR
+        + "%'";
+  }
 }

core/src/main/java/com/scalar/db/storage/jdbc/RdbEngineStrategy.java

@@ -305,4 +305,6 @@ default void setConnectionToReadOnly(Connection connection, boolean readOnly)
    */
   default void throwIfCrossPartitionScanOrderingOnBlobColumnNotSupported(
       ScanAll scanAll, TableMetadata metadata) {}
+
+  String getTableNamesInNamespaceSql(String namespace);
 }

core/src/test/java/com/scalar/db/storage/jdbc/JdbcAdminTest.java

@@ -1957,6 +1957,7 @@ public void dropNamespace_forSqlite_shouldDropNamespace() throws Exception {
     JdbcAdmin admin = createJdbcAdminFor(RdbEngine.SQLITE);
 
     Connection connection = mock(Connection.class);
+    PreparedStatement getTableNamesPrepStmt = mock(PreparedStatement.class);
     PreparedStatement deleteFromNamespaceTablePrepStmt = mock(PreparedStatement.class);
     Statement selectAllFromNamespaceTablePrepStmt = mock(Statement.class);
     Statement selectAllFromMetadataTablePrepStmt = mock(Statement.class);
@@ -1967,7 +1968,13 @@ public void dropNamespace_forSqlite_shouldDropNamespace() throws Exception {
             selectAllFromNamespaceTablePrepStmt,
             selectAllFromMetadataTablePrepStmt,
             dropNamespaceTableStmt);
-    when(connection.prepareStatement(anyString())).thenReturn(deleteFromNamespaceTablePrepStmt);
+    // Mock for getNamespaceTableNamesInternal() check - returns empty ResultSet (namespace is
+    // empty)
+    ResultSet emptyResultSet = mock(ResultSet.class);
+    when(emptyResultSet.next()).thenReturn(false);
+    when(getTableNamesPrepStmt.executeQuery()).thenReturn(emptyResultSet);
+    when(connection.prepareStatement(anyString()))
+        .thenReturn(getTableNamesPrepStmt, deleteFromNamespaceTablePrepStmt);
     when(dataSource.getConnection()).thenReturn(connection);
     // Only the metadata schema is left
     ResultSet resultSet1 =
@@ -2005,6 +2012,7 @@ private void dropNamespace_WithOnlyNamespaceSchemaLeftForX_shouldDropSchemaAndNa
 
     Connection connection = mock(Connection.class);
     Statement dropNamespaceStmt = mock(Statement.class);
+    PreparedStatement isNamespaceEmptyStatementMock = mock(PreparedStatement.class);
     PreparedStatement deleteFromNamespaceTablePrepStmt = mock(PreparedStatement.class);
     Statement selectAllFromNamespaceTablePrepStmt = mock(Statement.class);
     Statement selectAllFromMetadataTablePrepStmt = mock(Statement.class);
@@ -2018,7 +2026,12 @@ private void dropNamespace_WithOnlyNamespaceSchemaLeftForX_shouldDropSchemaAndNa
             selectAllFromMetadataTablePrepStmt,
             dropNamespaceTableStmt,
             dropMetadataSchemaStmt);
-    when(connection.prepareStatement(anyString())).thenReturn(deleteFromNamespaceTablePrepStmt);
+    // Mock for isNamespaceEmpty() check - returns empty ResultSet (namespace is empty)
+    ResultSet emptyResultSet = mock(ResultSet.class);
+    when(emptyResultSet.next()).thenReturn(false);
+    when(isNamespaceEmptyStatementMock.executeQuery()).thenReturn(emptyResultSet);
+    when(connection.prepareStatement(anyString()))
+        .thenReturn(isNamespaceEmptyStatementMock, deleteFromNamespaceTablePrepStmt);
     when(dataSource.getConnection()).thenReturn(connection);
     // Only the metadata schema is left
     ResultSet resultSet =
@@ -2116,6 +2129,7 @@ private void dropNamespace_WithOtherNamespaceLeftForX_shouldOnlyDropNamespace(
 
     Connection connection = mock(Connection.class);
     Statement dropNamespaceStatementMock = mock(Statement.class);
+    PreparedStatement isNamespaceEmptyStatementMock = mock(PreparedStatement.class);
     PreparedStatement deleteFromNamespaceTableMock = mock(PreparedStatement.class);
     Statement selectNamespaceStatementMock = mock(Statement.class);
     if (rdbEngine != RdbEngine.SQLITE) {
@@ -2124,7 +2138,12 @@ private void dropNamespace_WithOtherNamespaceLeftForX_shouldOnlyDropNamespace(
     } else {
       when(connection.createStatement()).thenReturn(selectNamespaceStatementMock);
     }
-    when(connection.prepareStatement(anyString())).thenReturn(deleteFromNamespaceTableMock);
+    // Mock for isNamespaceEmpty() check - returns empty ResultSet (namespace is empty)
+    ResultSet emptyResultSet = mock(ResultSet.class);
+    when(emptyResultSet.next()).thenReturn(false);
+    when(isNamespaceEmptyStatementMock.executeQuery()).thenReturn(emptyResultSet);
+    when(connection.prepareStatement(anyString()))
+        .thenReturn(isNamespaceEmptyStatementMock, deleteFromNamespaceTableMock);
     when(dataSource.getConnection()).thenReturn(connection);
     // Namespaces table contains other namespaces
     ResultSet resultSet =

integration-test/src/main/java/com/scalar/db/api/DistributedStorageAdminImportTableIntegrationTestBase.java

@@ -249,7 +249,7 @@ public void dropNamespace_ShouldNotDropNonScalarDBTables() throws Exception {
 
       // Act
       assertThatCode(() -> admin.dropNamespace(getNamespace()))
-          .isInstanceOf(ExecutionException.class);
+          .isInstanceOf(IllegalStateException.class);
 
       // Assert
       assertThat(admin.namespaceExists(getNamespace())).isTrue();