Fix JDBI imcompatiblity issues. Refactoring.

yoyama · Jan 12, 2022 · bf086cd · bf086cd
1 parent dc1dadf
commit bf086cd
Show file tree

Hide file tree

Showing 8 changed files with 23 additions and 24 deletions.
diff --git a/digdag-cli/src/main/java/io/digdag/cli/Migrate.java b/digdag-cli/src/main/java/io/digdag/cli/Migrate.java
@@ -8,7 +8,7 @@
 import io.digdag.core.database.DataSourceProvider;
 import io.digdag.core.database.DatabaseConfig;
 import io.digdag.core.database.DatabaseMigrator;
-import io.digdag.core.database.JdbiHelper;
+import io.digdag.core.database.DatabaseHelper;
 import io.digdag.core.database.migrate.Migration;
 import org.jdbi.v3.core.Jdbi;
 
@@ -34,7 +34,7 @@ public void main()
         checkArgs();
         DatabaseConfig dbConfig = DatabaseConfig.convertFrom(buildConfig());
         try (DataSourceProvider dsp = new DataSourceProvider(dbConfig)) {
-            Jdbi dbi = JdbiHelper.createJdbi(dsp.get());
+            Jdbi dbi = DatabaseHelper.createJdbi(dsp.get());
             DatabaseMigrator migrator = new DatabaseMigrator(dbi, dbConfig);
             switch (subCommand) {
                 case RUN:

diff --git a/...a/io/digdag/core/database/JdbiHelper.java → .../digdag/core/database/DatabaseHelper.java b/...a/io/digdag/core/database/JdbiHelper.java → .../digdag/core/database/DatabaseHelper.java
@@ -7,14 +7,18 @@
 
 import javax.sql.DataSource;
 
-public class JdbiHelper
+public class DatabaseHelper
 {
     public static Jdbi createJdbi(DataSource ds)
     {
         Jdbi jdbi = Jdbi.create(ds);
         jdbi.installPlugin(new SqlObjectPlugin());
-        jdbi.installPlugin(new PostgresPlugin());
-        jdbi.installPlugin(new H2DatabasePlugin());
+        if (ds.getClass().getCanonicalName().startsWith("org.h2")) {
+            jdbi.installPlugin(new H2DatabasePlugin());
+        }
+        else {
+            jdbi.installPlugin(new PostgresPlugin());
+        }
         return jdbi;
     }
 }
diff --git a/digdag-core/src/main/java/io/digdag/core/database/DatabaseModule.java b/digdag-core/src/main/java/io/digdag/core/database/DatabaseModule.java
@@ -51,7 +51,7 @@ public static class AutoMigrator
         public AutoMigrator(DataSource ds, DatabaseConfig config)
         {
             if (config.getAutoMigrate()) {
-                this.migrator = new DatabaseMigrator(JdbiHelper.createJdbi(ds), config);
+                this.migrator = new DatabaseMigrator(DatabaseHelper.createJdbi(ds), config);
             }
         }
 
@@ -80,7 +80,7 @@ public DbiProvider(DataSource ds, AutoMigrator migrator)
         @Override
         public Jdbi get()
         {
-            return JdbiHelper.createJdbi(ds);
+            return DatabaseHelper.createJdbi(ds);
         }
     }
 }
diff --git a/digdag-core/src/main/java/io/digdag/core/database/DatabaseProjectStoreManager.java b/digdag-core/src/main/java/io/digdag/core/database/DatabaseProjectStoreManager.java
@@ -410,7 +410,7 @@ public DatabaseProjectControlStore(Handle handle, int siteId)
         {
             this.handle = handle;
             this.siteId = siteId;
-            this.dao = handle.attach(Dao.class);
+            this.dao = handle.attach(dao(databaseType));
         }
 
         /**

diff --git a/digdag-core/src/main/java/io/digdag/core/database/DatabaseSessionStoreManager.java b/digdag-core/src/main/java/io/digdag/core/database/DatabaseSessionStoreManager.java
@@ -471,7 +471,7 @@ private <T> Optional<T> lockTask(long taskId, TaskLockAction<T> func, boolean if
         return transaction((handle, dao) -> {
             Long locked = ifNotLocked ? dao.lockTaskIfNotLocked(taskId) : dao.lockTask(taskId);
             if (locked != null) {
-                T result = func.call(new DatabaseTaskControlStore(handle, databaseType));
+                T result = func.call(new DatabaseTaskControlStore(handle));
                 return Optional.of(result);
             }
             return Optional.<T>absent();
@@ -501,7 +501,7 @@ private <T> Optional<T> lockTaskWithDetails(long taskId, TaskLockActionWithDetai
             if (locked != null) {
                 try {
                     StoredTask task = getTaskById(handle, taskId);
-                    T result = func.call(new DatabaseTaskControlStore(handle, databaseType), task);
+                    T result = func.call(new DatabaseTaskControlStore(handle), task);
                     return Optional.of(result);
                 }
                 catch (ResourceNotFoundException ex) {
@@ -783,7 +783,7 @@ public <T> T lockRootTask(long attemptId, TaskLockActionWithDetails<T> func)
                     dao.lockRootTask(attemptId),
                     "root task of attempt id=%d", attemptId);
             StoredTask task = getTaskById(handle, taskId);
-            T result = func.call(new DatabaseTaskControlStore(handle, databaseType), task);
+            T result = func.call(new DatabaseTaskControlStore(handle), task);
             return result;
         }
 
@@ -821,16 +821,11 @@ private class DatabaseTaskControlStore
         private final Handle handle;
         private final Dao dao;
 
-        public DatabaseTaskControlStore(Handle handle, String dbType)
+        public DatabaseTaskControlStore(Handle handle)
         {
             this.handle = handle;
             // JDBI3 does not accept Dao.class which is abstract.
-            if (dbType.equals("postgresql")) {
-                this.dao = handle.attach(PgDao.class);
-            }
-            else {
-                this.dao = handle.attach(H2Dao.class);
-            }
+            this.dao = handle.attach(dao(databaseType));
         }
 
 
@@ -958,7 +953,7 @@ public boolean copyInitialTasksForRetry(List<Long> recursiveChildrenIdList, Opti
                 return false;
             }
 
-            DatabaseTaskControlStore store = new DatabaseTaskControlStore(handle, databaseType);
+            DatabaseTaskControlStore store = new DatabaseTaskControlStore(handle);
             Map<Long, Long> oldIdToNewId = new HashMap<>();
             for (StoredTask task : tasks) {
                 Task newTask = Task.taskBuilder()
@@ -1562,7 +1557,7 @@ public <T> T insertRootTask(long attemptId, Task task, SessionBuilderAction<T> f
             long taskId = dao.insertTask(attemptId, task.getParentId().orNull(), task.getTaskType().get(), task.getState().get(), task.getStateFlags().get());  // tasks table don't have unique index
             dao.insertTaskDetails(taskId, task.getFullName(), task.getConfig().getLocal(), task.getConfig().getExport());
             dao.insertEmptyTaskStateDetails(taskId);
-            return func.call(new DatabaseTaskControlStore(handle, databaseType), taskId);
+            return func.call(new DatabaseTaskControlStore(handle), taskId);
         }
 
         @DigdagTimed(value = "dscst_", category = "db", appendMethodName = true)

diff --git a/digdag-core/src/main/java/io/digdag/core/database/ThreadLocalTransactionManager.java b/digdag-core/src/main/java/io/digdag/core/database/ThreadLocalTransactionManager.java
@@ -63,7 +63,7 @@ public Handle getHandle(ConfigMapper configMapper)
             }
 
             if (handle == null) {
-                Jdbi dbi = JdbiHelper.createJdbi(ds);
+                Jdbi dbi = DatabaseHelper.createJdbi(ds);
                 ConfigKeyListMapper cklm = new ConfigKeyListMapper();
                 dbi.registerRowMapper(new DatabaseProjectStoreManager.StoredProjectMapper(configMapper));
                 dbi.registerRowMapper(new DatabaseProjectStoreManager.StoredProjectWithRevisionMapper(configMapper));

diff --git a/digdag-core/src/test/java/io/digdag/core/database/DatabaseTestingUtils.java b/digdag-core/src/test/java/io/digdag/core/database/DatabaseTestingUtils.java
@@ -73,7 +73,7 @@ public static DatabaseFactory setupDatabase(boolean autoAutoCommit)
         DatabaseConfig config = getEnvironmentDatabaseConfig();
         DataSourceProvider dsp = new DataSourceProvider(config);
 
-        Jdbi dbi = JdbiHelper.createJdbi(dsp.get());
+        Jdbi dbi = DatabaseHelper.createJdbi(dsp.get());
         TransactionManager tm = new ThreadLocalTransactionManager(dsp.get(), autoAutoCommit);
         // FIXME
         new DatabaseMigrator(dbi, config).migrate();

diff --git a/digdag-tests/src/test/java/acceptance/MigrationIT.java b/digdag-tests/src/test/java/acceptance/MigrationIT.java
@@ -1,7 +1,7 @@
 package acceptance;
 
 import io.digdag.core.database.DatabaseMigrator;
-import io.digdag.core.database.JdbiHelper;
+import io.digdag.core.database.DatabaseHelper;
 import io.digdag.core.database.migrate.Migration;
 import io.digdag.core.database.migrate.MigrationContext;
 import io.digdag.core.database.migrate.Migration_20151204221156_CreateTables;
@@ -50,7 +50,7 @@ public void checkDatabaseUpgrade()
         try {
             server.setupDatabase();
             DataSource ds = server.getTestDBDataSource();
-            Jdbi dbi = JdbiHelper.createJdbi(ds);
+            Jdbi dbi = DatabaseHelper.createJdbi(ds);
             DatabaseMigrator migrator = new DatabaseMigrator(dbi, server.getRemoteTestDatabaseConfig());
             MigrationContext context = new MigrationContext(migrator.getDatabaseType());