Skip to content

Commit

Permalink
[apache#4882] improvement(core): Add some ut for MySQL and PostgreSQL…
Browse files Browse the repository at this point in the history
… storage backend. (apache#4898)

### What changes were proposed in this pull request?

Add some UTs to test MySQL and PostgreSQL storage backend.


### Why are the changes needed?

To improve code robustness.

Fix: apache#4882 

### Does this PR introduce _any_ user-facing change?

N/A.

### How was this patch tested?

N/A

---------

Co-authored-by: Jerry Shao <jerryshao@datastrato.com>
  • Loading branch information
yuqi1129 and jerryshao authored Sep 11, 2024
1 parent f48e9ad commit 228a1f5
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 79 deletions.
3 changes: 3 additions & 0 deletions core/build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -51,11 +51,14 @@ dependencies {
testCompileOnly(libs.lombok)

testImplementation(project(":integration-test-common", "testArtifacts"))
testImplementation(project(":server-common"))
testImplementation(project(":clients:client-java"))
testImplementation(libs.awaitility)
testImplementation(libs.junit.jupiter.api)
testImplementation(libs.junit.jupiter.params)
testImplementation(libs.mockito.core)
testImplementation(libs.mysql.driver)
testImplementation(libs.postgresql.driver)
testImplementation(libs.testcontainers)

testRuntimeOnly(libs.junit.jupiter.engine)
Expand Down
128 changes: 83 additions & 45 deletions core/src/test/java/org/apache/gravitino/storage/TestEntityStorage.java
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,7 @@

package org.apache.gravitino.storage;

import static org.apache.gravitino.Configs.DEFAULT_ENTITY_KV_STORE;
import static org.apache.gravitino.Configs.DEFAULT_ENTITY_RELATIONAL_STORE;
import static org.apache.gravitino.Configs.ENTITY_KV_ROCKSDB_BACKEND_PATH;
import static org.apache.gravitino.Configs.ENTITY_KV_STORE;
import static org.apache.gravitino.Configs.ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER;
import static org.apache.gravitino.Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD;
import static org.apache.gravitino.Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PATH;
Expand All @@ -32,7 +29,6 @@
import static org.apache.gravitino.Configs.ENTITY_STORE;
import static org.apache.gravitino.Configs.RELATIONAL_ENTITY_STORE;
import static org.apache.gravitino.Configs.STORE_DELETE_AFTER_TIME;
import static org.apache.gravitino.Configs.STORE_TRANSACTION_MAX_SKEW_TIME;
import static org.apache.gravitino.Configs.VERSION_RETENTION_COUNT;

import com.google.common.base.Preconditions;
Expand All @@ -51,6 +47,7 @@
import java.util.UUID;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.apache.gravitino.Catalog;
import org.apache.gravitino.Config;
import org.apache.gravitino.Configs;
Expand All @@ -68,6 +65,8 @@
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.exceptions.NonEmptyEntityException;
import org.apache.gravitino.file.Fileset;
import org.apache.gravitino.integration.test.container.ContainerSuite;
import org.apache.gravitino.integration.test.util.AbstractIT;
import org.apache.gravitino.meta.AuditInfo;
import org.apache.gravitino.meta.BaseMetalake;
import org.apache.gravitino.meta.CatalogEntity;
Expand All @@ -79,14 +78,25 @@
import org.apache.gravitino.meta.TableEntity;
import org.apache.gravitino.meta.TopicEntity;
import org.apache.gravitino.meta.UserEntity;
import org.apache.gravitino.storage.relational.converters.H2ExceptionConverter;
import org.apache.gravitino.storage.relational.converters.MySQLExceptionConverter;
import org.apache.gravitino.storage.relational.converters.PostgreSQLExceptionConverter;
import org.apache.gravitino.storage.relational.converters.SQLExceptionConverterFactory;
import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper;
import org.apache.ibatis.session.SqlSession;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;
import org.mockito.Mockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Tag("gravitino-docker-test")
public class TestEntityStorage {
private static final Logger LOG = LoggerFactory.getLogger(TestEntityStorage.class);

public static final String KV_STORE_PATH =
"/tmp/gravitino_kv_entityStore_" + UUID.randomUUID().toString().replace("-", "");

Expand All @@ -96,50 +106,72 @@ public class TestEntityStorage {
private static final String H2_FILE = DB_DIR + ".mv.db";

static Object[] storageProvider() {
return new Object[] {Configs.RELATIONAL_ENTITY_STORE};
return new Object[] {"h2", "mysql", "postgresql"};
}

@AfterEach
void closeSuit() throws IOException {
ContainerSuite.getInstance().close();
}

private void init(String type, Config config) {
Preconditions.checkArgument(StringUtils.isNotBlank(type));
if (type.equals(Configs.KV_STORE_KEY)) {
try {
FileUtils.deleteDirectory(FileUtils.getFile(KV_STORE_PATH));
} catch (Exception e) {
// Ignore
}
Mockito.when(config.get(ENTITY_STORE)).thenReturn("kv");
Mockito.when(config.get(ENTITY_KV_STORE)).thenReturn(DEFAULT_ENTITY_KV_STORE);
Mockito.when(config.get(Configs.ENTITY_SERDE)).thenReturn("proto");
Mockito.when(config.get(ENTITY_KV_ROCKSDB_BACKEND_PATH)).thenReturn(KV_STORE_PATH);

Assertions.assertEquals(KV_STORE_PATH, config.get(ENTITY_KV_ROCKSDB_BACKEND_PATH));
Mockito.when(config.get(STORE_TRANSACTION_MAX_SKEW_TIME)).thenReturn(1000L);
Mockito.when(config.get(STORE_DELETE_AFTER_TIME)).thenReturn(20 * 60 * 1000L);
Mockito.when(config.get(VERSION_RETENTION_COUNT)).thenReturn(1L);
} else if (type.equals(Configs.RELATIONAL_ENTITY_STORE)) {
File dir = new File(DB_DIR);
if (dir.exists() || !dir.isDirectory()) {
dir.delete();
File dir = new File(DB_DIR);
if (dir.exists() || !dir.isDirectory()) {
dir.delete();
}
dir.mkdirs();
Mockito.when(config.get(ENTITY_STORE)).thenReturn(RELATIONAL_ENTITY_STORE);
Mockito.when(config.get(ENTITY_RELATIONAL_STORE)).thenReturn(DEFAULT_ENTITY_RELATIONAL_STORE);
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_PATH)).thenReturn(DB_DIR);
Mockito.when(config.get(STORE_DELETE_AFTER_TIME)).thenReturn(20 * 60 * 1000L);
Mockito.when(config.get(VERSION_RETENTION_COUNT)).thenReturn(1L);

try {
if (type.equalsIgnoreCase("h2")) {
// The following properties are used to create the JDBC connection; they are just for test,
// in the real world, they will be set automatically by the configuration file if you set
// ENTITY_RELATIONAL_STOR as EMBEDDED_ENTITY_RELATIONAL_STORE.
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_URL))
.thenReturn(String.format("jdbc:h2:%s;DB_CLOSE_DELAY=-1;MODE=MYSQL", DB_DIR));
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_USER)).thenReturn("gravitino");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD)).thenReturn("gravitino");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER)).thenReturn("org.h2.Driver");

FieldUtils.writeStaticField(
SQLExceptionConverterFactory.class, "converter", new H2ExceptionConverter(), true);

} else if (type.equalsIgnoreCase("mysql")) {
String mysqlJdbcUrl = AbstractIT.startAndInitMySQLBackend();
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_URL)).thenReturn(mysqlJdbcUrl);
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_USER)).thenReturn("root");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD)).thenReturn("root");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER))
.thenReturn("com.mysql.cj.jdbc.Driver");

FieldUtils.writeStaticField(
SQLExceptionConverterFactory.class, "converter", new MySQLExceptionConverter(), true);

} else if (type.equalsIgnoreCase("postgresql")) {
String postgreSQLJdbcUrl = AbstractIT.startAndInitPGBackend();
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_URL)).thenReturn(postgreSQLJdbcUrl);
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_USER)).thenReturn("root");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD)).thenReturn("root");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER))
.thenReturn("org.postgresql.Driver");

FieldUtils.writeStaticField(
SQLExceptionConverterFactory.class,
"converter",
new PostgreSQLExceptionConverter(),
true);

} else {
throw new UnsupportedOperationException("Unsupported entity store type: " + type);
}
dir.mkdirs();
Mockito.when(config.get(ENTITY_STORE)).thenReturn(RELATIONAL_ENTITY_STORE);
Mockito.when(config.get(ENTITY_RELATIONAL_STORE)).thenReturn(DEFAULT_ENTITY_RELATIONAL_STORE);
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_PATH)).thenReturn(DB_DIR);

// The following properties are used to create the JDBC connection; they are just for test, in
// the real world,
// they will be set automatically by the configuration file if you set ENTITY_RELATIONAL_STORE
// as EMBEDDED_ENTITY_RELATIONAL_STORE.
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_URL))
.thenReturn(String.format("jdbc:h2:%s;DB_CLOSE_DELAY=-1;MODE=MYSQL", DB_DIR));
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_USER)).thenReturn("gravitino");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD)).thenReturn("gravitino");
Mockito.when(config.get(ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER)).thenReturn("org.h2.Driver");

Mockito.when(config.get(STORE_DELETE_AFTER_TIME)).thenReturn(20 * 60 * 1000L);
Mockito.when(config.get(VERSION_RETENTION_COUNT)).thenReturn(1L);
} else {
throw new UnsupportedOperationException("Unsupported entity store type: " + type);
} catch (Exception e) {
LOG.error("Failed to init entity store", e);
throw new RuntimeException(e);
}
}

Expand All @@ -151,14 +183,16 @@ private void destroy(String type) {
} catch (Exception e) {
// Ignore
}
} else if (type.equals(Configs.RELATIONAL_ENTITY_STORE)) {
} else if (type.equalsIgnoreCase("h2") || type.equalsIgnoreCase("mysql")) {
dropAllTables();
File dir = new File(DB_DIR);
if (dir.exists()) {
dir.delete();
}

FileUtils.deleteQuietly(new File(H2_FILE));
} else if (type.equalsIgnoreCase("postgresql")) {
// Do nothing
} else {
throw new UnsupportedOperationException("Unsupported entity store type: " + type);
}
Expand Down Expand Up @@ -876,6 +910,8 @@ void testSameNameUnderANameSpace(String type) throws IOException {
store.get(identifier, Entity.EntityType.TABLE, TableEntity.class);
store.get(identifier, Entity.EntityType.FILESET, FilesetEntity.class);
store.get(changedNameIdentifier, Entity.EntityType.TOPIC, TopicEntity.class);

destroy(type);
}
}

Expand Down Expand Up @@ -2104,7 +2140,7 @@ private void validateDeletedTable(EntityStore store) throws IOException {
@ParameterizedTest
@MethodSource("storageProvider")
void testOptimizedDeleteForKv(String type) throws IOException {
if ("relational".equalsIgnoreCase(type)) {
if (!"kv".equalsIgnoreCase(type)) {
return;
}

Expand Down Expand Up @@ -2182,6 +2218,8 @@ void testOptimizedDeleteForKv(String type) throws IOException {
Assertions.assertDoesNotThrow(
() ->
store.get(filesetEntity1.nameIdentifier(), EntityType.FILESET, FilesetEntity.class));

destroy(type);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -563,6 +563,17 @@ private static long[] cidrToRange(String cidr) throws Exception {
public void close() throws IOException {
try {
closer.close();
mySQLContainer = null;
mySQLVersion5Container = null;
hiveContainer = null;
hiveRangerContainer = null;
trinoContainer = null;
trinoITContainers = null;
rangerContainer = null;
kafkaContainer = null;
dorisContainer = null;
kerberosHiveContainer = null;
pgContainerMap.clear();
} catch (Exception e) {
LOG.error("Failed to close ContainerEnvironment", e);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@
import java.nio.file.Paths;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
Expand Down Expand Up @@ -161,21 +160,12 @@ protected static void downLoadJDBCDriver() throws IOException {
}
}

protected static void setPGBackend() throws SQLException {
String pgUrlWithoutSchema = POSTGRESQL_CONTAINER.getJdbcUrl(META_DATA);
customConfigs.put(Configs.ENTITY_STORE_KEY, "relational");
customConfigs.put(Configs.ENTITY_RELATIONAL_STORE_KEY, "JDBCBackend");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL_KEY, pgUrlWithoutSchema);
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER_KEY,
POSTGRESQL_CONTAINER.getDriverClassName(META_DATA));
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_USER_KEY, POSTGRESQL_CONTAINER.getUsername());
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD_KEY, POSTGRESQL_CONTAINER.getPassword());

LOG.info("PG URL: {}", pgUrlWithoutSchema);
public static String startAndInitPGBackend() {
META_DATA = PG_JDBC_BACKEND;
containerSuite.startPostgreSQLContainer(META_DATA);
POSTGRESQL_CONTAINER = containerSuite.getPostgreSQLContainer();

String pgUrlWithoutSchema = POSTGRESQL_CONTAINER.getJdbcUrl(META_DATA);
String randomSchemaName = RandomStringUtils.random(10, true, false);
// Connect to the PostgreSQL docker and create a schema
String currentExecuteSql = "";
Expand Down Expand Up @@ -216,18 +206,17 @@ protected static void setPGBackend() throws SQLException {

pgUrlWithoutSchema = pgUrlWithoutSchema + "?currentSchema=" + randomSchemaName;
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL_KEY, pgUrlWithoutSchema);

LOG.info("PG URL: {}", pgUrlWithoutSchema);
return pgUrlWithoutSchema;
}

private static void setMySQLBackend() {
String mysqlUrl = MYSQL_CONTAINER.getJdbcUrl(META_DATA);
customConfigs.put(Configs.ENTITY_STORE_KEY, "relational");
customConfigs.put(Configs.ENTITY_RELATIONAL_STORE_KEY, "JDBCBackend");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL_KEY, mysqlUrl);
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER_KEY, "com.mysql.cj.jdbc.Driver");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_USER_KEY, "root");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD_KEY, "root");
public static String startAndInitMySQLBackend() {
META_DATA = TestDatabaseName.MYSQL_JDBC_BACKEND;
containerSuite.startMySQLContainer(META_DATA);
MYSQL_CONTAINER = containerSuite.getMySQLContainer();

String mysqlUrl = MYSQL_CONTAINER.getJdbcUrl(META_DATA);
LOG.info("MySQL URL: {}", mysqlUrl);
// Connect to the mysql docker and create a databases
try (Connection connection =
Expand Down Expand Up @@ -255,6 +244,7 @@ private static void setMySQLBackend() {
for (String sql : initMySQLBackendSqls) {
statement.execute(sql);
}
return mysqlUrl;
} catch (Exception e) {
LOG.error("Failed to create database in mysql", e);
throw new RuntimeException(e);
Expand All @@ -279,18 +269,27 @@ public static void startIntegrationTest() throws Exception {

if ("MySQL".equalsIgnoreCase(System.getenv("jdbcBackend"))) {
// Start MySQL docker instance.
META_DATA = TestDatabaseName.MYSQL_JDBC_BACKEND;
containerSuite.startMySQLContainer(META_DATA);
MYSQL_CONTAINER = containerSuite.getMySQLContainer();

setMySQLBackend();
String jdbcURL = startAndInitMySQLBackend();
customConfigs.put(Configs.ENTITY_STORE_KEY, "relational");
customConfigs.put(Configs.ENTITY_RELATIONAL_STORE_KEY, "JDBCBackend");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL_KEY, jdbcURL);
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER_KEY, "com.mysql.cj.jdbc.Driver");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_USER_KEY, "root");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD_KEY, "root");
} else if ("PostgreSQL".equalsIgnoreCase(System.getenv("jdbcBackend"))) {
// Start PostgreSQL docker instance.
META_DATA = PG_JDBC_BACKEND;
containerSuite.startPostgreSQLContainer(META_DATA);
POSTGRESQL_CONTAINER = containerSuite.getPostgreSQLContainer();

setPGBackend();
String pgJdbcUrl = startAndInitPGBackend();
customConfigs.put(Configs.ENTITY_STORE_KEY, "relational");
customConfigs.put(Configs.ENTITY_RELATIONAL_STORE_KEY, "JDBCBackend");
customConfigs.put(Configs.ENTITY_RELATIONAL_JDBC_BACKEND_URL_KEY, pgJdbcUrl);
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_DRIVER_KEY,
POSTGRESQL_CONTAINER.getDriverClassName(META_DATA));
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_USER_KEY, POSTGRESQL_CONTAINER.getUsername());
customConfigs.put(
Configs.ENTITY_RELATIONAL_JDBC_BACKEND_PASSWORD_KEY, POSTGRESQL_CONTAINER.getPassword());
}

File baseDir = new File(System.getProperty("java.io.tmpdir"));
Expand Down

0 comments on commit 228a1f5

Please sign in to comment.