feat: improve sql audit store test (#722)

davidfrigolet · web-flow · commit f69dc742c4c0 · 2025-10-31T09:40:23.000Z
diff --git a/community/flamingock-auditstore-sql/src/main/java/io/flamingock/community/sql/internal/SqlAuditorDialectHelper.java b/community/flamingock-auditstore-sql/src/main/java/io/flamingock/community/sql/internal/SqlAuditorDialectHelper.java
@@ -138,28 +138,33 @@ public String getCreateTableSqlString(String tableName) {
                                 ")'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -955 THEN RAISE; END IF; END;", tableName, getAutoIncrementType(), getClobType(), getBigIntType(), getClobType(), getBooleanType(), getBooleanType());
             case DB2:
                 return String.format(
-                        "CREATE TABLE %s (" +
-                                "id %s PRIMARY KEY, " +
-                                "execution_id VARCHAR(255), " +
-                                "stage_id VARCHAR(255), " +
-                                "task_id VARCHAR(255), " +
-                                "author VARCHAR(255), " +
-                                "created_at TIMESTAMP DEFAULT CURRENT TIMESTAMP, " +
-                                "state VARCHAR(255), " +
-                                "class_name VARCHAR(255), " +
-                                "method_name VARCHAR(255), " +
-                                "metadata %s, " +
-                                "execution_millis %s, " +
-                                "execution_hostname VARCHAR(255), " +
-                                "error_trace %s, " +
-                                "type VARCHAR(50), " +
-                                "tx_type VARCHAR(50), " +
-                                "target_system_id VARCHAR(255), " +
-                                "order_col VARCHAR(50), " +
-                                "recovery_strategy VARCHAR(50), " +
-                                "transaction_flag %s, " +
-                                "system_change %s" +
-                                ")", tableName, getAutoIncrementType(), getClobType(), getBigIntType(), getClobType(), getBooleanType(), getBooleanType());
+                        "BEGIN\n" +
+                                "  DECLARE CONTINUE HANDLER FOR SQLSTATE '42710' BEGIN END;\n" +
+                                "  EXECUTE IMMEDIATE 'CREATE TABLE %s (" +
+                                "id BIGINT NOT NULL GENERATED ALWAYS AS IDENTITY (START WITH 1, INCREMENT BY 1) PRIMARY KEY," +
+                                "execution_id VARCHAR(255)," +
+                                "stage_id VARCHAR(255)," +
+                                "task_id VARCHAR(255)," +
+                                "author VARCHAR(255)," +
+                                "created_at TIMESTAMP," +
+                                "state VARCHAR(32)," +
+                                "class_name VARCHAR(255)," +
+                                "method_name VARCHAR(255)," +
+                                "metadata CLOB," +
+                                "execution_millis BIGINT," +
+                                "execution_hostname VARCHAR(255)," +
+                                "error_trace CLOB," +
+                                "type VARCHAR(50)," +
+                                "tx_type VARCHAR(50)," +
+                                "target_system_id VARCHAR(255)," +
+                                "order_col VARCHAR(50)," +
+                                "recovery_strategy VARCHAR(50)," +
+                                "transaction_flag SMALLINT," +
+                                "system_change SMALLINT" +
+                                ")';\n" +
+                                "END",
+                        tableName);
+
             case SQLITE:
                 return String.format(
                         "CREATE TABLE IF NOT EXISTS %s (" +
diff --git a/community/flamingock-auditstore-sql/src/main/java/io/flamingock/community/sql/internal/SqlLockDialectHelper.java b/community/flamingock-auditstore-sql/src/main/java/io/flamingock/community/sql/internal/SqlLockDialectHelper.java
@@ -78,12 +78,14 @@ public String getCreateTableSqlString(String tableName) {
                                 ")'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -955 THEN RAISE; END IF; END;", tableName);
             case DB2:
                 return String.format(
-                        "CREATE TABLE %s (" +
-                                "\"key\" VARCHAR(255) PRIMARY KEY," +
-                                "status VARCHAR(32)," +
-                                "owner VARCHAR(255)," +
-                                "expires_at TIMESTAMP" +
-                                ")", tableName);
+                        "BEGIN " +
+                                "DECLARE CONTINUE HANDLER FOR SQLSTATE '42710' BEGIN END; " +
+                                "EXECUTE IMMEDIATE 'CREATE TABLE %s (" +
+                                "lock_key VARCHAR(255) NOT NULL PRIMARY KEY, " +
+                                "status VARCHAR(32), " +
+                                "owner VARCHAR(255), " +
+                                "expires_at TIMESTAMP)'; " +
+                                "END", tableName);
             default:
                 throw new UnsupportedOperationException("Dialect not supported for CREATE TABLE: " + sqlDialect.name());
         }
@@ -93,6 +95,9 @@ public String getSelectLockSqlString(String tableName) {
         switch (sqlDialect) {
             case POSTGRESQL:
                 return String.format("SELECT \"key\", status, owner, expires_at FROM %s WHERE \"key\" = ?", tableName);
+            case DB2:
+                // Select lock_key as the first column (getLockEntry expects rs.getString(1) to be the key)
+                return String.format("SELECT lock_key, status, owner, expires_at FROM %s WHERE lock_key = ?", tableName);
             case SQLSERVER:
             case SYBASE:
                 return String.format("SELECT [key], status, owner, expires_at FROM %s WITH (UPDLOCK, ROWLOCK) WHERE [key] = ?", tableName);
@@ -145,12 +150,13 @@ public String getInsertOrUpdateLockSqlString(String tableName) {
                         "MERGE INTO %s (`key`, status, owner, expires_at) KEY (`key`) VALUES (?, ?, ?, ?)",
                         tableName);
             case DB2:
+                // Use a VALUES-derived table and a target alias for DB2 to avoid parsing issues
                 return String.format(
-                        "MERGE INTO %s USING (SELECT ? AS \"key\", ? AS status, ? AS owner, ? AS expires_at FROM SYSIBM.SYSDUMMY1) AS src " +
-                                "ON (%s.\"key\" = src.\"key\") " +
+                        "MERGE INTO %s tgt USING (VALUES (?, ?, ?, ?)) src(lock_key, status, owner, expires_at) " +
+                                "ON (tgt.lock_key = src.lock_key) " +
                                 "WHEN MATCHED THEN UPDATE SET status = src.status, owner = src.owner, expires_at = src.expires_at " +
-                                "WHEN NOT MATCHED THEN INSERT (\"key\", status, owner, expires_at) VALUES (src.\"key\", src.status, src.owner, src.expires_at)",
-                        tableName, tableName);
+                                "WHEN NOT MATCHED THEN INSERT (lock_key, status, owner, expires_at) VALUES (src.lock_key, src.status, src.owner, src.expires_at)",
+                        tableName);
             case FIREBIRD:
                 return String.format(
                         "UPDATE OR INSERT INTO %s (`key`, status, owner, expires_at) VALUES (?, ?, ?, ?) MATCHING (`key`)",
@@ -164,12 +170,41 @@ public String getDeleteLockSqlString(String tableName) {
         if (Objects.requireNonNull(sqlDialect) == SqlDialect.POSTGRESQL) {
             return String.format("DELETE FROM %s WHERE \"key\" = ?", tableName);
         }
+        if (sqlDialect == SqlDialect.DB2) {
+            return String.format("DELETE FROM %s WHERE lock_key = ?", tableName);
+        }
         return String.format("DELETE FROM %s WHERE `key` = ?", tableName);
     }
 
     public void upsertLockEntry(Connection conn, String tableName, String key, String owner, LocalDateTime expiresAt) throws SQLException {
         String sql = getInsertOrUpdateLockSqlString(tableName);
 
+        if (sqlDialect == SqlDialect.DB2) {
+            // UPDATE first
+            try (PreparedStatement update = conn.prepareStatement(
+                    "UPDATE " + tableName + " SET owner = ?, expires_at = ? WHERE lock_key = ?")) {
+                update.setString(1, owner);
+                update.setTimestamp(2, Timestamp.valueOf(expiresAt));
+                update.setString(3, key);
+                int updated = update.executeUpdate();
+                if (updated > 0) {
+                    return;
+                }
+            }
+
+            // If no row updated, try INSERT
+            try (PreparedStatement insert = conn.prepareStatement(
+                    "INSERT INTO " + tableName + " (lock_key, status, owner, expires_at) VALUES (?, ?, ?, ?)")) {
+                insert.setString(1, key);
+                // Use "LOCKED" string to avoid using a non-existing enum constant (previous "ACQUIRED" caused failures)
+                insert.setString(2, LockStatus.LOCK_HELD.name());
+                insert.setString(3, owner);
+                insert.setTimestamp(4, Timestamp.valueOf(expiresAt));
+                insert.executeUpdate();
+                return;
+            }
+        }
+
         if (getSqlDialect() == SqlDialect.SQLSERVER || getSqlDialect() == SqlDialect.SYBASE) {
             // For SQL Server/Sybase, use Statement and format SQL
             try (Statement stmt = conn.createStatement()) {
@@ -186,6 +221,7 @@ public void upsertLockEntry(Connection conn, String tableName, String key, Strin
             }
         } else {
             try (PreparedStatement ps = conn.prepareStatement(sql)) {
+                // For DB2 we use lock_key but callers pass key as first parameter - that's correct
                 ps.setString(1, key);
                 ps.setString(2, LockStatus.LOCK_HELD.name());
                 ps.setString(3, owner);
diff --git a/community/flamingock-auditstore-sql/src/main/java/io/flamingock/community/sql/internal/SqlLockService.java b/community/flamingock-auditstore-sql/src/main/java/io/flamingock/community/sql/internal/SqlLockService.java
@@ -136,7 +136,7 @@ public void releaseLock(LockKey lockKey, RunnerId owner) {
         String keyStr = lockKey.toString();
         try (Connection conn = dataSource.getConnection();
              PreparedStatement ps = conn.prepareStatement(
-                     "SELECT owner FROM " + lockRepositoryName + " WHERE `key` = ?")) {
+                     dialectHelper.getSelectLockSqlString(lockRepositoryName))) {
             ps.setString(1, keyStr);
             try (ResultSet rs = ps.executeQuery()) {
                 if (rs.next()) {
diff --git a/community/flamingock-auditstore-sql/src/test/java/io/flamingock/community/sql/SharedSqlContainers.java b/community/flamingock-auditstore-sql/src/test/java/io/flamingock/community/sql/SharedSqlContainers.java
@@ -0,0 +1,111 @@
+/*
+ * Copyright 2025 Flamingock (https://www.flamingock.io)
+ *
+ * Licensed 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 io.flamingock.community.sql;
+
+import org.testcontainers.containers.*;
+import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.containers.wait.strategy.WaitAllStrategy;
+import org.testcontainers.utility.DockerImageName;
+
+import java.time.Duration;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.sql.DataSource;
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+
+public final class SharedSqlContainers {
+
+    private static final ConcurrentHashMap<String, JdbcDatabaseContainer<?>> CONTAINERS = new ConcurrentHashMap<>();
+
+    private SharedSqlContainers() { }
+
+    public static JdbcDatabaseContainer<?> getContainer(String dialectName) {
+        boolean isCi = System.getenv("CI") != null || System.getenv("GITHUB_ACTIONS") != null;
+        return CONTAINERS.computeIfAbsent(dialectName, key -> createContainerInternal(key, isCi));
+    }
+
+    private static JdbcDatabaseContainer<?> createContainerInternal(String dialectName, boolean isCi) {
+        switch (dialectName) {
+            case "mysql": {
+                MySQLContainer<?> c = new MySQLContainer<>("mysql:8.0")
+                        .withDatabaseName("testdb")
+                        .withUsername("testuser")
+                        .withPassword("testpass");
+                if (!isCi) c.withReuse(true);
+                return c;
+            }
+            case "sqlserver": {
+                MSSQLServerContainer<?> c = new MSSQLServerContainer<>("mcr.microsoft.com/mssql/server:2019-CU18-ubuntu-20.04")
+                        .acceptLicense()
+                        .withPassword("TestPass123!");
+                if (!isCi) c.withReuse(true);
+                return c;
+            }
+            case "oracle": {
+                OracleContainer c = new OracleContainer(
+                        DockerImageName.parse("gvenzl/oracle-free:23-slim-faststart")
+                                .asCompatibleSubstituteFor("gvenzl/oracle-xe")) {
+                    @Override
+                    public String getDatabaseName() {
+                        return "FREEPDB1";
+                    }
+                }
+                        .withPassword("oracle123")
+                        .withSharedMemorySize(1073741824L)
+                        .withStartupTimeout(Duration.ofMinutes(20))
+                        .waitingFor(new WaitAllStrategy()
+                                .withStrategy(Wait.forListeningPort())
+                                .withStrategy(Wait.forLogMessage(".*DATABASE IS READY TO USE.*\\n", 1))
+                        )
+                        .withEnv("ORACLE_CHARACTERSET", "AL32UTF8");
+                if (!isCi) c.withReuse(true);
+                return c;
+            }
+            case "postgresql": {
+                PostgreSQLContainer<?> c = new PostgreSQLContainer<>(DockerImageName.parse("postgres:15"))
+                        .withDatabaseName("testdb")
+                        .withUsername("test")
+                        .withPassword("test");
+                if (!isCi) c.withReuse(true);
+                return c;
+            }
+            case "mariadb": {
+                MariaDBContainer<?> c = new MariaDBContainer<>("mariadb:11.3.2")
+                        .withDatabaseName("testdb")
+                        .withUsername("testuser")
+                        .withPassword("testpass");
+                if (!isCi) c.withReuse(true);
+                return c;
+            }
+            default:
+                throw new IllegalArgumentException("Unsupported dialect: " + dialectName);
+        }
+    }
+
+    public static void stopAll() {
+        CONTAINERS.values().forEach(JdbcDatabaseContainer::stop);
+        CONTAINERS.clear();
+    }
+
+    public static DataSource createDataSource(JdbcDatabaseContainer<?> container) {
+        HikariConfig config = new HikariConfig();
+        config.setJdbcUrl(container.getJdbcUrl());
+        config.setUsername(container.getUsername());
+        config.setPassword(container.getPassword());
+        config.setDriverClassName(container.getDriverClassName());
+        return new HikariDataSource(config);
+    }
+}
diff --git a/community/flamingock-auditstore-sql/src/test/java/io/flamingock/community/sql/SqlAuditStoreTest.java b/community/flamingock-auditstore-sql/src/test/java/io/flamingock/community/sql/SqlAuditStoreTest.java
diff --git a/community/flamingock-auditstore-sql/src/test/java/io/flamingock/community/sql/SqlAuditTestHelper.java b/community/flamingock-auditstore-sql/src/test/java/io/flamingock/community/sql/SqlAuditTestHelper.java
