From 2ac6d8418f578790ddfc453f99d5d646cd9645f1 Mon Sep 17 00:00:00 2001
From: Nicklas Wallgren <nicklas.wallgren@gmail.com>
Date: Fri, 7 Feb 2025 09:00:09 +0100
Subject: [PATCH] feat: Add option to provide a custom DataSource

---
 .../autoconfigure/DataSourceSupplier.java     | 13 +++++++++++++
 .../DbSchedulerAutoConfiguration.java         | 19 +++++++++++++------
 2 files changed, 26 insertions(+), 6 deletions(-)
 create mode 100644 db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DataSourceSupplier.java

diff --git a/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DataSourceSupplier.java b/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DataSourceSupplier.java
new file mode 100644
index 00000000..a1029ac6
--- /dev/null
+++ b/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DataSourceSupplier.java
@@ -0,0 +1,13 @@
+package com.github.kagkarlsson.scheduler.boot.autoconfigure;
+
+import javax.sql.DataSource;
+
+@FunctionalInterface
+public interface DataSourceSupplier {
+  /**
+   * Supplies an instance of {@link DataSource}.
+   *
+   * @return The {@link DataSource} to be used by the scheduler.
+   */
+  DataSource dataSource();
+}
diff --git a/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DbSchedulerAutoConfiguration.java b/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DbSchedulerAutoConfiguration.java
index 3fdb7822..9c3ae868 100644
--- a/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DbSchedulerAutoConfiguration.java
+++ b/db-scheduler-boot-starter/src/main/java/com/github/kagkarlsson/scheduler/boot/autoconfigure/DbSchedulerAutoConfiguration.java
@@ -68,22 +68,18 @@ public class DbSchedulerAutoConfiguration {
   private static final Predicate<Task<?>> shouldBeStarted = task -> task instanceof OnStartup;
 
   private final DbSchedulerProperties config;
-  private final DataSource existingDataSource;
   private final List<Task<?>> configuredTasks;
   private final List<SchedulerListener> schedulerListeners;
   private final List<ExecutionInterceptor> executionInterceptors;
 
   public DbSchedulerAutoConfiguration(
       DbSchedulerProperties dbSchedulerProperties,
-      DataSource dataSource,
       List<Task<?>> configuredTasks,
       List<SchedulerListener> schedulerListeners,
       List<ExecutionInterceptor> executionInterceptors) {
     this.config =
         Objects.requireNonNull(
             dbSchedulerProperties, "Can't configure db-scheduler without required configuration");
-    this.existingDataSource =
-        Objects.requireNonNull(dataSource, "An existing javax.sql.DataSource is required");
     this.configuredTasks =
         Objects.requireNonNull(configuredTasks, "At least one Task must be configured");
     this.schedulerListeners = schedulerListeners;
@@ -105,15 +101,24 @@ StatsRegistry noopStatsRegistry() {
     return StatsRegistry.NOOP;
   }
 
+  @ConditionalOnMissingBean(DataSourceSupplier.class)
+  @Bean
+  public DataSourceSupplier dataSourceSupplier(final DataSource existingDataSource) {
+    return () -> existingDataSource;
+  }
+
   @ConditionalOnBean(DataSource.class)
   @ConditionalOnMissingBean
   @DependsOnDatabaseInitialization
   @Bean(destroyMethod = "stop")
-  public Scheduler scheduler(DbSchedulerCustomizer customizer, StatsRegistry registry) {
+  public Scheduler scheduler(
+      DbSchedulerCustomizer customizer,
+      StatsRegistry registry,
+      DataSourceSupplier dataSourceSupplier) {
     log.info("Creating db-scheduler using tasks from Spring context: {}", configuredTasks);
 
     // Ensure that we are using a transactional aware data source
-    DataSource transactionalDataSource = configureDataSource(existingDataSource);
+    DataSource transactionalDataSource = configureDataSource(dataSourceSupplier.dataSource());
 
     // Instantiate a new builder
     final SchedulerBuilder builder =
@@ -217,6 +222,8 @@ static LazyInitializationExcludeFilter eagerDbSchedulerStarter() {
   }
 
   private static DataSource configureDataSource(DataSource existingDataSource) {
+    Objects.requireNonNull(existingDataSource, "An existing javax.sql.DataSource is required");
+
     if (existingDataSource instanceof TransactionAwareDataSourceProxy) {
       log.debug("Using an already transaction aware DataSource");
       return existingDataSource;