#37 Simplify MySQL support

bucket4j · Feb 9, 2022 · 5743f9c · 5743f9c
1 parent f6c602c
commit 5743f9c
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 142 deletions.
diff --git a/bucket4j-mysql/src/main/java/io/github/bucket4j/mysql/MySQLProxyConfigurationBuilder.java b/bucket4j-mysql/src/main/java/io/github/bucket4j/mysql/MySQLProxyConfigurationBuilder.java
@@ -26,7 +26,7 @@ public static MySQLProxyConfigurationBuilder builder() {
      *                                                 By default, under the hood uses {@link ClientSideConfig#getDefault}
      * @return {@link MySQLProxyConfigurationBuilder}
      */
-    public MySQLProxyConfigurationBuilder addClientSideConfig(ClientSideConfig clientSideConfig) {
+    public MySQLProxyConfigurationBuilder withClientSideConfig(ClientSideConfig clientSideConfig) {
         this.clientSideConfig = clientSideConfig;
         return this;
     }
@@ -36,7 +36,7 @@ public MySQLProxyConfigurationBuilder addClientSideConfig(ClientSideConfig clien
      *                                                 By default, under the hood uses {@link BucketTableSettings#getDefault}
      * @return {@link MySQLProxyConfigurationBuilder}
      */
-    public MySQLProxyConfigurationBuilder addTableSettings(BucketTableSettings tableSettings) {
+    public MySQLProxyConfigurationBuilder withTableSettings(BucketTableSettings tableSettings) {
         this.tableSettings = tableSettings;
         return this;
     }

diff --git a/bucket4j-mysql/src/main/java/io/github/bucket4j/mysql/MySQLProxyManager.java b/bucket4j-mysql/src/main/java/io/github/bucket4j/mysql/MySQLProxyManager.java
@@ -7,6 +7,7 @@
 import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.text.MessageFormat;
 import java.util.Objects;
@@ -24,6 +25,9 @@ public class MySQLProxyManager extends AbstractLockBasedProxyManager<Long> {
     private final DataSource dataSource;
     private final MySQLProxyConfiguration configuration;
     private final String removeSqlQuery;
+    private final String updateSqlQuery;
+    private final String insertSqlQuery;
+    private final String selectSqlQuery;
 
     /**
      *
@@ -34,15 +38,118 @@ public MySQLProxyManager(MySQLProxyConfiguration configuration) {
         this.dataSource = Objects.requireNonNull(configuration.getDataSource());
         this.configuration = configuration;
         this.removeSqlQuery = MessageFormat.format("DELETE FROM {0} WHERE {1} = ?", configuration.getTableName(), configuration.getIdName());
+        updateSqlQuery = MessageFormat.format("UPDATE {0} SET {1}=? WHERE {2}=?", configuration.getTableName(), configuration.getStateName(), configuration.getIdName());
+        insertSqlQuery = MessageFormat.format("INSERT IGNORE INTO {0}({1}, {2}) VALUES(?, null)",
+                configuration.getTableName(), configuration.getIdName(), configuration.getStateName(), configuration.getIdName());
+        selectSqlQuery = MessageFormat.format("SELECT {0} FROM {1} WHERE {2} = ? FOR UPDATE", configuration.getStateName(), configuration.getTableName(), configuration.getIdName());
     }
 
     @Override
     protected LockBasedTransaction allocateTransaction(Long key) {
+        Connection connection;
         try {
-            return new MySQLSelectForUpdateLockBasedTransaction(key, configuration, dataSource.getConnection());
+            connection = dataSource.getConnection();
         } catch (SQLException e) {
             throw new BucketExceptions.BucketExecutionException(e);
         }
+
+        return new LockBasedTransaction() {
+            @Override
+            public void begin() {
+                try {
+                    connection.setAutoCommit(false);
+                    PreparedStatement sth = connection.prepareStatement("SELECT GET_LOCK(?, ?)");
+                    sth.setLong(1, key);
+                    sth.setInt(2, 2);
+                    sth.executeQuery();
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+            @Override
+            public byte[] lockAndGet() {
+                try {
+                    try (PreparedStatement selectStatement = connection.prepareStatement(selectSqlQuery)) {
+                        selectStatement.setLong(1, key);
+                        try (ResultSet rs = selectStatement.executeQuery()) {
+                            if (rs.next()) {
+                                byte[] bucketStateBeforeTransaction = rs.getBytes(configuration.getStateName());
+                                if (bucketStateBeforeTransaction != null) {
+                                    return bucketStateBeforeTransaction;
+                                } else {
+                                    return null;
+                                }
+                            }
+                        }
+                    }
+                    try (PreparedStatement insertStatement = connection.prepareStatement(insertSqlQuery)) {
+                        insertStatement.setLong(1, key);
+                        insertStatement.executeUpdate();
+                    }
+                    try (PreparedStatement selectStatement = connection.prepareStatement(selectSqlQuery)) {
+                        selectStatement.setLong(1, key);
+                        try (ResultSet rs = selectStatement.executeQuery()) {
+                            if (!rs.next()) {
+                                throw new IllegalStateException("Something unexpected happens, it needs to read the MySQL manual");
+                            }
+                            return null;
+                        }
+                    }
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+            @Override
+            public void update(byte[] data) {
+                try {
+                    try (PreparedStatement updateStatement = connection.prepareStatement(updateSqlQuery)) {
+                        updateStatement.setBytes(1, data);
+                        updateStatement.setLong(2, key);
+                        updateStatement.executeUpdate();
+                    }
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+            @Override
+            public void release() {
+                try {
+                    connection.close();
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+            @Override
+            public void create(byte[] data) {
+                update(data);
+            }
+            @Override
+            public void rollback() {
+                try {
+                    connection.rollback();
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+            @Override
+            public void commit() {
+                try {
+                    connection.commit();
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+            @Override
+            public void unlock() {
+                try {
+                    PreparedStatement sth = connection.prepareStatement("SELECT RELEASE_LOCK(?)");
+                    sth.setLong(1, key);
+                    sth.executeQuery();
+                } catch (SQLException e) {
+                    throw new BucketExceptions.BucketExecutionException(e);
+                }
+            }
+        };
     }
 
     @Override
@@ -57,4 +164,5 @@ public void removeProxy(Long key) {
             throw new BucketExceptions.BucketExecutionException(e);
         }
     }
+
 }
diff --git a/...ysql/src/main/java/io/github/bucket4j/mysql/MySQLSelectForUpdateLockBasedTransaction.java b/...ysql/src/main/java/io/github/bucket4j/mysql/MySQLSelectForUpdateLockBasedTransaction.java
diff --git a/.../src/test/java/io/github/bucket4j/mysql/MySQLSelectForUpdateLockBasedTransactionTest.java b/.../src/test/java/io/github/bucket4j/mysql/MySQLSelectForUpdateLockBasedTransactionTest.java
@@ -36,8 +36,8 @@ public static void initializeInstance() throws SQLException {
             }
         }
         MySQLProxyConfiguration configuration = MySQLProxyConfigurationBuilder.builder()
-                .addClientSideConfig(ClientSideConfig.getDefault())
-                .addTableSettings(tableSettings)
+                .withClientSideConfig(ClientSideConfig.getDefault())
+                .withTableSettings(tableSettings)
                 .build(dataSource);
         proxyManager = new MySQLProxyManager(configuration);
     }