From 4c650874da206b93d42a6f41cf828bf115858e96 Mon Sep 17 00:00:00 2001
From: Librazy <im@librazy.org>
Date: Mon, 14 Jan 2019 13:53:54 +0800
Subject: [PATCH] Fix createTable within transaction

---
 .../database/provider/SQLiteDatabase.java     | 67 +++++++++++++++----
 1 file changed, 54 insertions(+), 13 deletions(-)

diff --git a/src/main/java/cat/nyaa/nyaacore/database/provider/SQLiteDatabase.java b/src/main/java/cat/nyaa/nyaacore/database/provider/SQLiteDatabase.java
index 2a89a264..8abcf167 100644
--- a/src/main/java/cat/nyaa/nyaacore/database/provider/SQLiteDatabase.java
+++ b/src/main/java/cat/nyaa/nyaacore/database/provider/SQLiteDatabase.java
@@ -19,6 +19,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.logging.Level;
@@ -32,7 +33,7 @@ public class SQLiteDatabase extends BaseDatabase {
     public static Function<Plugin, Consumer<Runnable>> executorSupplier = (plugin) -> (runnable) -> Bukkit.getScheduler().runTaskAsynchronously(plugin, runnable);
     public static Function<Plugin, Logger> loggerSupplier = Plugin::getLogger;
     private final Semaphore mainConnLock = new Semaphore(1);
-
+    private final AtomicBoolean mainConnInTransaction = new AtomicBoolean(false);
     private static FinalizableReferenceQueue frq = new FinalizableReferenceQueue();
     private static final ConcurrentMap<Reference<?>, Semaphore> references = Maps.newConcurrentMap();
 
@@ -94,8 +95,9 @@ protected Object clone() throws CloneNotSupportedException {
     }
 
     @Override
-    @SuppressWarnings("rawtypes")
     public void createTable(Class<?> cls) {
+        Validate.notNull(cls);
+        if (createdTableClasses.contains(cls)) return;
         try {
             if (!mainConnLock.tryAcquire(10, TimeUnit.SECONDS)) {
                 throw new IllegalStateException();
@@ -104,21 +106,25 @@ public void createTable(Class<?> cls) {
             throw new RuntimeException(e);
         }
         try {
-            Validate.notNull(cls);
-            if (createdTableClasses.contains(cls)) return;
-            TableStructure ts = TableStructure.fromClass(cls);
-            String sql = ts.getCreateTableSQL("sqlite");
-            try (Statement smt = getConnection().createStatement()) {
-                smt.executeUpdate(sql);
-                createdTableClasses.add(cls);
-            } catch (SQLException ex) {
-                throw new RuntimeException(sql, ex);
-            }
+            createTable(cls, getConnection());
         } finally {
             mainConnLock.release();
         }
     }
 
+    private void createTable(Class<?> cls, Connection conn) {
+        Validate.notNull(cls);
+        if (createdTableClasses.contains(cls)) return;
+        TableStructure ts = TableStructure.fromClass(cls);
+        String sql = ts.getCreateTableSQL("sqlite");
+        try (Statement smt = conn.createStatement()) {
+            smt.executeUpdate(sql);
+            createdTableClasses.add(cls);
+        } catch (SQLException ex) {
+            throw new RuntimeException(sql, ex);
+        }
+    }
+
     @Override
     public Plugin getPlugin() {
         return plugin;
@@ -129,34 +135,68 @@ public void beginTransaction() {
         try {
             mainConnLock.acquireUninterruptibly();
             super.beginTransaction();
+            mainConnInTransaction.set(true);
         } catch (Throwable e) {
             mainConnLock.release();
+            mainConnInTransaction.set(false);
             throw e;
         }
     }
 
     @Override
     public void commitTransaction() {
+        if (!mainConnInTransaction.get()) {
+            throw new IllegalStateException();
+        }
         try {
             super.commitTransaction();
         } finally {
             mainConnLock.release();
+            mainConnInTransaction.set(false);
         }
     }
 
     @Override
     public void rollbackTransaction() {
+        if (!mainConnInTransaction.get()) {
+            throw new IllegalStateException();
+        }
         try {
             super.rollbackTransaction();
         } finally {
             mainConnLock.release();
+            mainConnInTransaction.set(false);
         }
     }
 
+    /**
+     * Return the SynchronizedQuery object for specified table class.
+     *
+     * @return SynchronizedQuery object
+     */
+    @Override
+    public <T> SynchronizedQuery.NonTransactionalQuery<T> query(Class<T> tableClass) {
+        if (mainConnInTransaction.get()) {
+            createTable(tableClass, getConnection());
+        } else {
+            createTable(tableClass);
+        }
+        return new SynchronizedQuery.NonTransactionalQuery<T>(tableClass, this.getConnection()) {
+            @Override
+            public T selectUniqueForUpdate() {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public void close() {
+
+            }
+        };
+    }
+
     @Override
     public <T> SynchronizedQuery.TransactionalQuery<T> queryTransactional(Class<T> tableClass) {
         Connection conn;
-        createTable(tableClass);
         try {
             if (!mainConnLock.tryAcquire(10, TimeUnit.SECONDS)) {
                 throw new IllegalStateException();
@@ -167,6 +207,7 @@ public <T> SynchronizedQuery.TransactionalQuery<T> queryTransactional(Class<T> t
         try {
             conn = getConnection();
             conn.setAutoCommit(false);
+            createTable(tableClass, conn);
         } catch (Throwable ex) {
             mainConnLock.release();
             throw new RuntimeException(ex);