Copilot commented on code in PR #11066: URL: https://github.com/apache/gravitino/pull/11066#discussion_r3240762240
########## plugins/idp-basic/src/test/java/org/apache/gravitino/idp/basic/storage/relational/BackendTestExtension.java: ########## @@ -0,0 +1,246 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF 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 org.apache.gravitino.idp.basic.storage.relational; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Stream; +import org.apache.gravitino.Config; +import org.apache.gravitino.Configs; +import org.apache.gravitino.integration.test.util.BaseIT; +import org.apache.gravitino.storage.relational.JDBCBackend; +import org.junit.jupiter.api.extension.AfterAllCallback; +import org.junit.jupiter.api.extension.BeforeAllCallback; +import org.junit.jupiter.api.extension.BeforeEachCallback; +import org.junit.jupiter.api.extension.Extension; +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.api.extension.TestTemplateInvocationContext; +import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider; + +public class BackendTestExtension + implements TestTemplateInvocationContextProvider, BeforeAllCallback, AfterAllCallback { + + private static final String DOCKER_TEST_FLAG = "dockerTest"; + private static final ExtensionContext.Namespace NAMESPACE = + ExtensionContext.Namespace.create(BackendTestExtension.class); + private static final String STORE_KEY = "BACKEND_MAP"; + + @Override + public void beforeAll(ExtensionContext context) { + context.getStore(NAMESPACE).put(STORE_KEY, new ConcurrentHashMap<String, BackendResource>()); + } Review Comment: Because `BackendResource`s are created and stored per `ExtensionContext` (i.e. per test class), the MySQL/PostgreSQL Testcontainers backends will be started and stopped for every docker-tagged test class instead of once per JVM. Additionally, `createBackendResource` calls `jdbcBackend.close()` on a fresh `JDBCBackend` to reset the static `SqlSessionFactoryHelper` singleton before re-initializing it, which means the singleton is torn down and re-initialized between test classes and prevents running tests of different backends in parallel. Consider hoisting the backend cache to a `BeforeAllCallback`-registered root context store (e.g. `ExtensionContext.Store.CloseableResource` on the `GLOBAL`/root namespace) so containers and the SqlSessionFactory are initialized once and properly closed at JVM shutdown. ########## plugins/idp-basic/src/main/java/org/apache/gravitino/idp/basic/storage/relational/mapper/IdpUserMetaSQLProviderFactory.java: ########## @@ -0,0 +1,112 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF 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 org.apache.gravitino.idp.basic.storage.relational.mapper; + +import com.google.common.collect.ImmutableMap; +import java.util.List; +import java.util.Map; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.base.IdpUserMetaBaseSQLProvider; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.h2.IdpUserMetaH2Provider; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.mysql.IdpUserMetaMySQLProvider; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.postgresql.IdpUserMetaPostgreSQLProvider; +import org.apache.gravitino.idp.basic.storage.relational.po.IdpUserPO; +import org.apache.gravitino.storage.relational.JDBCBackend.JDBCBackendType; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.apache.ibatis.annotations.Param; + +public class IdpUserMetaSQLProviderFactory { + private static final Map<JDBCBackendType, IdpUserMetaBaseSQLProvider> + IDP_USER_META_SQL_PROVIDER_MAP = + ImmutableMap.of( + JDBCBackendType.MYSQL, new IdpUserMetaMySQLProvider(), + JDBCBackendType.H2, new IdpUserMetaH2Provider(), + JDBCBackendType.POSTGRESQL, new IdpUserMetaPostgreSQLProvider()); + + public static IdpUserMetaBaseSQLProvider getProvider() { + String databaseId = + SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getDatabaseId(); + + return getProvider(databaseId); + } + + static IdpUserMetaBaseSQLProvider getProvider(String databaseId) { + if (databaseId == null) { + throw new IllegalStateException( + "MyBatis databaseId is not configured for IdP user SQL providers."); + } + + JDBCBackendType jdbcBackendType; + try { + jdbcBackendType = JDBCBackendType.fromString(databaseId); + } catch (IllegalArgumentException e) { + throw new IllegalStateException( + String.format( + "Unsupported IdP user SQL provider databaseId: %s, supported backends: %s", + databaseId, IDP_USER_META_SQL_PROVIDER_MAP.keySet()), + e); + } + + return getProvider(jdbcBackendType, databaseId, IDP_USER_META_SQL_PROVIDER_MAP); + } + + static IdpUserMetaBaseSQLProvider getProvider( + JDBCBackendType jdbcBackendType, + String databaseId, + Map<JDBCBackendType, IdpUserMetaBaseSQLProvider> providerMap) { + IdpUserMetaBaseSQLProvider provider = providerMap.get(jdbcBackendType); Review Comment: This package-private 3-argument overload exists only to let `TestIdpUserMetaSQLProviderFactoryFailure` inject an empty map and exercise the "no provider registered" branch. That branch is unreachable in production because `IDP_USER_META_SQL_PROVIDER_MAP` contains an entry for every `JDBCBackendType` returned by `JDBCBackendType.fromString` here. Consider either (a) removing this overload and the corresponding test (the unreachable branch can be dropped), or (b) collapsing the lookup so a missing entry is treated the same as an unsupported `databaseId` in the 1-arg method, keeping the public surface and tests aligned with real failure modes. ########## plugins/idp-basic/src/test/java/org/apache/gravitino/idp/basic/storage/relational/BackendTestExtension.java: ########## @@ -0,0 +1,246 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF 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 org.apache.gravitino.idp.basic.storage.relational; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.stream.Stream; +import org.apache.gravitino.Config; +import org.apache.gravitino.Configs; +import org.apache.gravitino.integration.test.util.BaseIT; +import org.apache.gravitino.storage.relational.JDBCBackend; +import org.junit.jupiter.api.extension.AfterAllCallback; +import org.junit.jupiter.api.extension.BeforeAllCallback; +import org.junit.jupiter.api.extension.BeforeEachCallback; +import org.junit.jupiter.api.extension.Extension; +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.api.extension.TestTemplateInvocationContext; +import org.junit.jupiter.api.extension.TestTemplateInvocationContextProvider; + +public class BackendTestExtension + implements TestTemplateInvocationContextProvider, BeforeAllCallback, AfterAllCallback { + + private static final String DOCKER_TEST_FLAG = "dockerTest"; + private static final ExtensionContext.Namespace NAMESPACE = + ExtensionContext.Namespace.create(BackendTestExtension.class); + private static final String STORE_KEY = "BACKEND_MAP"; + + @Override + public void beforeAll(ExtensionContext context) { + context.getStore(NAMESPACE).put(STORE_KEY, new ConcurrentHashMap<String, BackendResource>()); + } + + @Override + @SuppressWarnings("unchecked") + public void afterAll(ExtensionContext context) throws Exception { + ConcurrentHashMap<String, BackendResource> map = + (ConcurrentHashMap<String, BackendResource>) context.getStore(NAMESPACE).get(STORE_KEY); + if (map != null) { + for (BackendResource backendResource : map.values()) { + backendResource.close(); + } + map.clear(); + } + } + + @Override + public boolean supportsTestTemplate(ExtensionContext context) { + return TestJDBCBackend.class.isAssignableFrom(context.getRequiredTestClass()); + } + + @Override + public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts( + ExtensionContext context) { + return resolveBackends(context.getRequiredTestClass()).stream() + .map(BackendInvocationContext::new); + } + + private List<String> resolveBackends(Class<?> testClass) { + BackendTypes backendTypes = findBackendTypes(testClass); + if (backendTypes != null) { + return List.of(backendTypes.value()); + } + + List<String> backendsToTest = new ArrayList<>(); + backendsToTest.add("h2"); + if ("true".equalsIgnoreCase(System.getenv(DOCKER_TEST_FLAG))) { Review Comment: The rest of the Gravitino test infrastructure typically gates docker-based tests via the `dockerTest` Gradle/JUnit system property (and/or the `gravitino-docker-test` JUnit tag), not via an environment variable. Reading `System.getenv("dockerTest")` here will silently skip MySQL/PostgreSQL invocations in normal `./gradlew test -PdockerTest=...` runs. Consider switching to `System.getProperty(DOCKER_TEST_FLAG)` (and/or relying solely on the `@Tag("gravitino-docker-test")` already present on the docker test classes) to stay consistent with the rest of the project. ########## plugins/idp-basic/src/test/java/org/apache/gravitino/idp/basic/storage/relational/TestJDBCBackend.java: ########## @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF 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 org.apache.gravitino.idp.basic.storage.relational; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import org.apache.gravitino.integration.test.util.CloseContainerExtension; +import org.apache.gravitino.integration.test.util.PrintFuncNameExtension; +import org.apache.gravitino.storage.relational.JDBCBackend; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.apache.ibatis.session.SqlSession; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.ExtendWith; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@ExtendWith({ + BackendTestExtension.class, + PrintFuncNameExtension.class, + CloseContainerExtension.class +}) +public abstract class TestJDBCBackend { + protected String backendType; + protected JDBCBackend backend; + + public void setBackendType(String backendType) { + this.backendType = backendType; + } + + public void setBackend(JDBCBackend backend) { + this.backend = backend; + } + + @BeforeEach + public void init() throws SQLException { + truncateAllTables(); + } + + protected long queryLongValue(String table, String column, String idColumn, long idValue) { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection()) { + String query = "SELECT " + column + " FROM " + table + " WHERE " + idColumn + " = ?"; + try (PreparedStatement statement = connection.prepareStatement(query)) { + statement.setLong(1, idValue); + try (ResultSet resultSet = statement.executeQuery()) { + assertTrue(resultSet.next()); + return resultSet.getLong(1); + } + } + } + } catch (SQLException e) { + throw new RuntimeException("Query " + column + " from " + table + " failed", e); + } + } + + protected int countRows(String table, String idColumn, long idValue) { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection()) { + String query = "SELECT COUNT(1) FROM " + table + " WHERE " + idColumn + " = ?"; + try (PreparedStatement statement = connection.prepareStatement(query)) { + statement.setLong(1, idValue); + try (ResultSet resultSet = statement.executeQuery()) { + assertTrue(resultSet.next()); + return resultSet.getInt(1); + } + } + } + } catch (SQLException e) { + throw new RuntimeException("Count rows from " + table + " failed", e); + } + } + + protected String currentJdbcUrl() { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection()) { + DatabaseMetaData metaData = connection.getMetaData(); + return metaData.getURL(); + } + } catch (SQLException e) { + throw new RuntimeException("Get current JDBC URL failed", e); + } + } + + private void truncateAllTables() throws SQLException { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection(); + Statement statement = connection.createStatement()) { + if ("postgresql".equalsIgnoreCase(backendType)) { + truncateAllTablesForPostgreSQL(connection); + } else { + List<String> tableList = new ArrayList<>(); + try (ResultSet rs = statement.executeQuery("SHOW TABLES")) { + while (rs.next()) { + tableList.add(rs.getString(1)); + } + } + for (String table : tableList) { + statement.execute("TRUNCATE TABLE " + table); + } Review Comment: `SHOW TABLES` returns table names without schema qualification and includes any backend-managed tables (e.g. MyBatis/Flyway/Liquibase metadata tables if present). On H2, `SHOW TABLES` also returns a `TABLE_SCHEMA` column whose value (depending on MySQL compatibility mode) may cause `TRUNCATE TABLE <name>` to target an unintended object, and for MySQL it indiscriminately truncates every table in the connected schema. Consider restricting to the application schema explicitly (e.g. query `information_schema.tables` with `table_schema = DATABASE()` / `current_schema()` and filter to the tables this plugin actually owns, such as `idp_user_meta`) to make truncation deterministic and safe to run in shared databases. ########## plugins/idp-basic/src/main/java/org/apache/gravitino/idp/basic/storage/relational/mapper/IdpUserMetaSQLProviderFactory.java: ########## @@ -0,0 +1,112 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF 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 org.apache.gravitino.idp.basic.storage.relational.mapper; + +import com.google.common.collect.ImmutableMap; +import java.util.List; +import java.util.Map; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.base.IdpUserMetaBaseSQLProvider; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.h2.IdpUserMetaH2Provider; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.mysql.IdpUserMetaMySQLProvider; +import org.apache.gravitino.idp.basic.storage.relational.mapper.provider.postgresql.IdpUserMetaPostgreSQLProvider; +import org.apache.gravitino.idp.basic.storage.relational.po.IdpUserPO; +import org.apache.gravitino.storage.relational.JDBCBackend.JDBCBackendType; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.apache.ibatis.annotations.Param; + +public class IdpUserMetaSQLProviderFactory { + private static final Map<JDBCBackendType, IdpUserMetaBaseSQLProvider> + IDP_USER_META_SQL_PROVIDER_MAP = + ImmutableMap.of( + JDBCBackendType.MYSQL, new IdpUserMetaMySQLProvider(), + JDBCBackendType.H2, new IdpUserMetaH2Provider(), + JDBCBackendType.POSTGRESQL, new IdpUserMetaPostgreSQLProvider()); + + public static IdpUserMetaBaseSQLProvider getProvider() { + String databaseId = + SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getDatabaseId(); + + return getProvider(databaseId); + } + + static IdpUserMetaBaseSQLProvider getProvider(String databaseId) { + if (databaseId == null) { + throw new IllegalStateException( + "MyBatis databaseId is not configured for IdP user SQL providers."); + } + + JDBCBackendType jdbcBackendType; + try { + jdbcBackendType = JDBCBackendType.fromString(databaseId); + } catch (IllegalArgumentException e) { + throw new IllegalStateException( + String.format( + "Unsupported IdP user SQL provider databaseId: %s, supported backends: %s", + databaseId, IDP_USER_META_SQL_PROVIDER_MAP.keySet()), + e); + } + + return getProvider(jdbcBackendType, databaseId, IDP_USER_META_SQL_PROVIDER_MAP); + } + + static IdpUserMetaBaseSQLProvider getProvider( + JDBCBackendType jdbcBackendType, + String databaseId, + Map<JDBCBackendType, IdpUserMetaBaseSQLProvider> providerMap) { + IdpUserMetaBaseSQLProvider provider = providerMap.get(jdbcBackendType); + if (provider == null) { + throw new IllegalStateException( + String.format( + "No IdP user SQL provider registered for backend %s (databaseId: %s)", + jdbcBackendType, databaseId)); + } + + return provider; + } + + public static String selectIdpUser(@Param("username") String username) { + return getProvider().selectIdpUser(username); + } Review Comment: The `@Param` annotations on these static SQL-builder methods (lines 87–110) have no effect — MyBatis only reads `@Param` from mapper interface methods, not from `*Provider` methods that return SQL strings. They're already present (and meaningful) on the mapper methods in `IdpUserMetaMapper`. Removing them here would avoid suggesting that these annotations participate in parameter binding. ########## plugins/idp-basic/src/test/java/org/apache/gravitino/idp/basic/storage/relational/TestJDBCBackend.java: ########## @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF 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 org.apache.gravitino.idp.basic.storage.relational; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import org.apache.gravitino.integration.test.util.CloseContainerExtension; +import org.apache.gravitino.integration.test.util.PrintFuncNameExtension; +import org.apache.gravitino.storage.relational.JDBCBackend; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.apache.ibatis.session.SqlSession; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.extension.ExtendWith; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@ExtendWith({ + BackendTestExtension.class, + PrintFuncNameExtension.class, + CloseContainerExtension.class +}) +public abstract class TestJDBCBackend { + protected String backendType; + protected JDBCBackend backend; + + public void setBackendType(String backendType) { + this.backendType = backendType; + } + + public void setBackend(JDBCBackend backend) { + this.backend = backend; + } + + @BeforeEach + public void init() throws SQLException { + truncateAllTables(); + } + + protected long queryLongValue(String table, String column, String idColumn, long idValue) { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection()) { + String query = "SELECT " + column + " FROM " + table + " WHERE " + idColumn + " = ?"; + try (PreparedStatement statement = connection.prepareStatement(query)) { + statement.setLong(1, idValue); + try (ResultSet resultSet = statement.executeQuery()) { + assertTrue(resultSet.next()); + return resultSet.getLong(1); + } + } + } + } catch (SQLException e) { + throw new RuntimeException("Query " + column + " from " + table + " failed", e); + } + } + + protected int countRows(String table, String idColumn, long idValue) { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection()) { + String query = "SELECT COUNT(1) FROM " + table + " WHERE " + idColumn + " = ?"; + try (PreparedStatement statement = connection.prepareStatement(query)) { + statement.setLong(1, idValue); + try (ResultSet resultSet = statement.executeQuery()) { + assertTrue(resultSet.next()); + return resultSet.getInt(1); + } + } + } + } catch (SQLException e) { + throw new RuntimeException("Count rows from " + table + " failed", e); + } + } + + protected String currentJdbcUrl() { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection()) { + DatabaseMetaData metaData = connection.getMetaData(); + return metaData.getURL(); + } + } catch (SQLException e) { + throw new RuntimeException("Get current JDBC URL failed", e); + } + } + + private void truncateAllTables() throws SQLException { + try (SqlSession sqlSession = + SqlSessionFactoryHelper.getInstance().getSqlSessionFactory().openSession(true)) { + try (Connection connection = sqlSession.getConnection(); + Statement statement = connection.createStatement()) { + if ("postgresql".equalsIgnoreCase(backendType)) { + truncateAllTablesForPostgreSQL(connection); + } else { + List<String> tableList = new ArrayList<>(); + try (ResultSet rs = statement.executeQuery("SHOW TABLES")) { + while (rs.next()) { + tableList.add(rs.getString(1)); + } + } + for (String table : tableList) { + statement.execute("TRUNCATE TABLE " + table); + } + } + } + } + } + + private void truncateAllTablesForPostgreSQL(Connection connection) throws SQLException { + List<String> tableList = new ArrayList<>(); + try (Statement statement = connection.createStatement()) { + String query = + "SELECT table_name FROM information_schema.tables WHERE table_schema = current_schema()"; + try (ResultSet rs = statement.executeQuery(query)) { + while (rs.next()) { + tableList.add(rs.getString(1)); + } + } + + if (tableList.isEmpty()) { + return; + } + + StringBuilder pgTruncateCommand = new StringBuilder("DO $$ BEGIN\n"); + for (String table : tableList) { + pgTruncateCommand.append( + String.format("TRUNCATE TABLE %s RESTART IDENTITY CASCADE;", table)); + } + pgTruncateCommand.append("END $$;"); + statement.execute(pgTruncateCommand.toString()); + } + } Review Comment: Table names are interpolated directly into a `DO $$ ... $$` block without quoting. Although the names come from `information_schema.tables` and are unlikely to be hostile, any identifier containing uppercase letters, reserved words, or special characters (which PostgreSQL allows when quoted) will produce invalid SQL or unexpectedly target a different relation. Wrap each identifier with `"` (or use `quote_ident(...)` in plpgsql) when composing the statement. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
