From 20c953f6fbdf1d94320403bf52e4382bb38b53c9 Mon Sep 17 00:00:00 2001 From: "John K. Watson" Date: Thu, 8 Feb 2024 15:59:30 -0800 Subject: [PATCH] feat: add an option to have a simpler JDBI execution threadpool --- .../java/ai/verta/modeldb/common/db/JdbiUtils.java | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/backend/common/src/main/java/ai/verta/modeldb/common/db/JdbiUtils.java b/backend/common/src/main/java/ai/verta/modeldb/common/db/JdbiUtils.java index 63a4a82070..541d17a693 100644 --- a/backend/common/src/main/java/ai/verta/modeldb/common/db/JdbiUtils.java +++ b/backend/common/src/main/java/ai/verta/modeldb/common/db/JdbiUtils.java @@ -7,6 +7,8 @@ import ai.verta.modeldb.common.futures.InternalJdbi; import com.zaxxer.hikari.HikariDataSource; import com.zaxxer.hikari.metrics.prometheus.PrometheusMetricsTrackerFactory; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import javax.sql.DataSource; import org.jdbi.v3.core.Jdbi; @@ -19,6 +21,15 @@ public static FutureJdbi initializeFutureJdbi( return new FutureJdbi(jdbi, dbExecutor); } + public static FutureJdbi initializeFutureJdbiWithFixedThreadpool( + DataSource dataSource, int poolSize) { + var jdbi = new InternalJdbi(Jdbi.create(dataSource)); + + ExecutorService executor = Executors.newFixedThreadPool(poolSize); + var dbExecutor = FutureExecutor.makeCompatibleExecutor(executor, "jdbi"); + return new FutureJdbi(jdbi, dbExecutor); + } + public static DataSource initializeDataSource(DatabaseConfig databaseConfig, String poolName) { final var hikariDataSource = new HikariDataSource(); final var dbUrl = RdbConfig.buildDatabaseConnectionString(databaseConfig.getRdbConfiguration());