Merge pull request #461 from DependencyTrack/jdbi

Add JDBI integration

pom.xml

@@ -92,6 +92,7 @@
         <lib.jackson.version>2.15.2</lib.jackson.version>
         <lib.jackson-databind.version>2.15.2</lib.jackson-databind.version>
         <lib.jaxb.runtime.version>2.3.6</lib.jaxb.runtime.version>
+        <lib.jdbi.version>3.42.0</lib.jdbi.version>
         <lib.json-unit.version>3.2.2</lib.json-unit.version>
         <!--
           Overriding logstash-logback-encoder version as a temporary workaround until
@@ -217,6 +218,16 @@
             <artifactId>cvss-calculator</artifactId>
             <version>${lib.cvss-calculator.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.jdbi</groupId>
+            <artifactId>jdbi3-core</artifactId>
+            <version>${lib.jdbi.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.jdbi</groupId>
+            <artifactId>jdbi3-sqlobject</artifactId>
+            <version>${lib.jdbi.version}</version>
+        </dependency>
         <!-- OWASP Risk Rating calculator -->
         <dependency>
             <groupId>us.springett</groupId>
@@ -544,6 +555,24 @@
                 <filtering>false</filtering>
             </testResource>
         </testResources>
+        <pluginManagement>
+            <plugins>
+                <plugin>
+                    <groupId>org.apache.maven.plugins</groupId>
+                    <artifactId>maven-compiler-plugin</artifactId>
+                    <configuration>
+                        <source>17</source>
+                        <target>17</target>
+                        <compilerArgs>
+                            <arg>-Xlint:all</arg>
+                            <arg>-Xlint:-processing</arg>
+                            <arg>-Xlint:-serial</arg>
+                            <arg>-parameters</arg>
+                        </compilerArgs>
+                    </configuration>
+                </plugin>
+            </plugins>
+        </pluginManagement>
         <plugins>
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>

src/main/java/org/dependencytrack/persistence/jdbi/JdbiFactory.java

@@ -0,0 +1,114 @@
+package org.dependencytrack.persistence.jdbi;
+
+import org.datanucleus.api.jdo.JDOPersistenceManagerFactory;
+import org.datanucleus.store.connection.ConnectionManagerImpl;
+import org.datanucleus.store.rdbms.ConnectionFactoryImpl;
+import org.datanucleus.store.rdbms.RDBMSStoreManager;
+import org.dependencytrack.persistence.QueryManager;
+import org.jdbi.v3.core.Jdbi;
+import org.jdbi.v3.sqlobject.SqlObjectPlugin;
+
+import javax.jdo.PersistenceManager;
+import javax.jdo.PersistenceManagerFactory;
+import javax.sql.DataSource;
+import java.sql.Connection;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.apache.commons.lang3.reflect.FieldUtils.readField;
+
+public class JdbiFactory {
+
+    private static final AtomicReference<GlobalInstanceHolder> GLOBAL_INSTANCE_HOLDER = new AtomicReference<>();
+
+    /**
+     * Get a global {@link Jdbi} instance, initializing it if it hasn't been initialized before.
+     * <p>
+     * The global instance will use {@link Connection}s from the primary {@link DataSource}
+     * of the given {@link QueryManager}'s {@link PersistenceManagerFactory}.
+     * <p>
+     * Usage of the global instance should be preferred to make the best possible use of JDBI's
+     * internal caching mechanisms. However, this instance can't participate in transactions
+     * initiated by JDO (via {@link QueryManager} or {@link PersistenceManager}).
+     * <p>
+     * If {@link Jdbi} usage in an active JDO {@link javax.jdo.Transaction} is desired,
+     * use {@link #localJdbi(QueryManager)} instead, which will use the same {@link Connection}
+     * as the provided {@link QueryManager}.
+     *
+     * @param qm The {@link QueryManager} to determine the {@link DataSource} from
+     * @return The global {@link Jdbi} instance
+     */
+    public static Jdbi jdbi(final QueryManager qm) {
+        return jdbi(qm.getPersistenceManager());
+    }
+
+    private static Jdbi jdbi(final PersistenceManager pm) {
+        return GLOBAL_INSTANCE_HOLDER
+                .updateAndGet(previous -> {
+                    if (previous == null || previous.pmf != pm.getPersistenceManagerFactory()) {
+                        // The PMF reference does not usually change, unless it has been recreated,
+                        // or multiple PMFs exist in the same application. The latter is not the case
+                        // for Dependency-Track, and the former only happens during test execution,
+                        // where each test (re-)creates the PMF.
+                        final Jdbi jdbi = createFromPmf(pm.getPersistenceManagerFactory());
+                        return new GlobalInstanceHolder(jdbi, pm.getPersistenceManagerFactory());
+                    }
+
+                    return previous;
+                })
+                .jdbi();
+    }
+
+    /**
+     * Create a new local {@link Jdbi} instance.
+     * <p>
+     * The instance will use the same {@link Connection} used by the given {@link QueryManager},
+     * allowing it to participate in {@link javax.jdo.Transaction}s initiated by {@code qm}.
+     * <p>
+     * Because using local {@link Jdbi} instances has a high performance impact (e.g. due to ineffective caching),
+     * this method will throw if {@code qm} is not participating in an active {@link javax.jdo.Transaction}
+     * already.
+     * <p>
+     * Just like {@link QueryManager} itself, {@link Jdbi} instances created by this method are <em>not</em>
+     * thread safe!
+     *
+     * @param qm The {@link QueryManager} to use the underlying {@link Connection} of
+     * @return A new {@link Jdbi} instance
+     * @throws IllegalStateException When the given {@link QueryManager} is not participating
+     *                               in an active {@link javax.jdo.Transaction}
+     */
+    public static Jdbi localJdbi(final QueryManager qm) {
+        return localJdbi(qm.getPersistenceManager());
+    }
+
+    private static Jdbi localJdbi(final PersistenceManager pm) {
+        if (!pm.currentTransaction().isActive()) {
+            throw new IllegalStateException("""
+                    Local JDBI instances must not be used outside of an active JDO transaction. \
+                    Use the global instance instead if combining JDBI with JDO transactions is not needed.""");
+        }
+
+        return Jdbi.create(new JdoConnectionFactory(pm));
+    }
+
+    private record GlobalInstanceHolder(Jdbi jdbi, PersistenceManagerFactory pmf) {
+    }
+
+    private static Jdbi createFromPmf(final PersistenceManagerFactory pmf) {
+        try {
+            if (pmf instanceof final JDOPersistenceManagerFactory jdoPmf
+                    && jdoPmf.getNucleusContext().getStoreManager() instanceof final RDBMSStoreManager storeManager
+                    && storeManager.getConnectionManager() instanceof final ConnectionManagerImpl connectionManager
+                    && readField(connectionManager, "primaryConnectionFactory", true) instanceof ConnectionFactoryImpl connectionFactory
+                    && readField(connectionFactory, "dataSource", true) instanceof final DataSource dataSource) {
+                return Jdbi
+                        .create(dataSource)
+                        .installPlugin(new SqlObjectPlugin());
+            }
+        } catch (IllegalAccessException e) {
+            throw new IllegalStateException("Failed to access datasource of PMF via reflection", e);
+        }
+
+        throw new IllegalStateException("Failed to access primary datasource of PMF");
+    }
+
+}

src/main/java/org/dependencytrack/persistence/jdbi/JdoConnectionFactory.java

@@ -0,0 +1,36 @@
+package org.dependencytrack.persistence.jdbi;
+
+import net.jcip.annotations.NotThreadSafe;
+import org.jdbi.v3.core.ConnectionFactory;
+
+import javax.jdo.PersistenceManager;
+import javax.jdo.datastore.JDOConnection;
+import java.sql.Connection;
+
+@NotThreadSafe
+class JdoConnectionFactory implements ConnectionFactory {
+
+    private final PersistenceManager pm;
+    private JDOConnection jdoConnection;
+
+    JdoConnectionFactory(final PersistenceManager pm) {
+        this.pm = pm;
+    }
+
+    @Override
+    public Connection openConnection() {
+        if (jdoConnection != null) {
+            throw new IllegalStateException("A JDO connection is already open");
+        }
+
+        jdoConnection = pm.getDataStoreConnection();
+        return (Connection) jdoConnection.getNativeConnection();
+    }
+
+    @Override
+    public void closeConnection(final Connection conn) {
+        jdoConnection.close();
+        jdoConnection = null;
+    }
+
+}

src/test/java/org/dependencytrack/persistence/jdbi/JdbiFactoryTest.java

@@ -0,0 +1,111 @@
+package org.dependencytrack.persistence.jdbi;
+
+import alpine.server.persistence.PersistenceManagerFactory;
+import org.dependencytrack.PersistenceCapableTest;
+import org.dependencytrack.model.Component;
+import org.dependencytrack.model.Project;
+import org.dependencytrack.persistence.QueryManager;
+import org.jdbi.v3.core.Jdbi;
+import org.junit.Test;
+
+import java.util.Optional;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+
+public class JdbiFactoryTest extends PersistenceCapableTest {
+
+    @Test
+    public void testGlobalInstance() {
+        final Jdbi jdbi = JdbiFactory.jdbi(qm);
+
+        // Issue a test query to ensure the JDBI instance is functional.
+        final Integer queryResult = jdbi.withHandle(handle ->
+                handle.createQuery("SELECT 666").mapTo(Integer.class).one());
+        assertThat(queryResult).isEqualTo(666);
+
+        // Ensure that the same JDBI instance is returned even when passing
+        // a different QueryManager instance. Because the underlying PMF
+        // did not change, the global JDBI instance must remain untouched.
+        try (final var otherQm = new QueryManager()) {
+            assertThat(JdbiFactory.jdbi(otherQm)).isEqualTo(jdbi);
+        }
+    }
+
+    @Test
+    public void testGlobalInstanceWithJdoTransaction() {
+        qm.runInTransaction(() -> {
+            // Create a new project.
+            final var project = new Project();
+            project.setName("acme-app");
+            project.setVersion("1.0.0");
+            qm.getPersistenceManager().makePersistent(project);
+
+            // Query for the created project, despite its creation not having been committed yet.
+            // Because the global JDBI instance uses a different connection than the QueryManager,
+            // it won't be able to see the yet-uncommitted change.
+            final Optional<String> projectName = JdbiFactory.jdbi(qm).withHandle(handle ->
+                    handle.createQuery("SELECT \"NAME\" FROM \"PROJECT\"").mapTo(String.class).findFirst());
+            assertThat(projectName).isNotPresent();
+        });
+    }
+
+    @Test
+    public void testGlobalInstanceWhenPmfChanges() {
+        final Jdbi jdbi = JdbiFactory.jdbi(qm);
+
+        // Close the PMF and ensure that the JDBI instance is no longer usable.
+        PersistenceManagerFactory.tearDown();
+        assertThatExceptionOfType(IllegalStateException.class)
+                .isThrownBy(() -> jdbi.withHandle(handle ->
+                        handle.createQuery("SELECT 666").mapTo(Integer.class).one()))
+                .withMessage("Pool not open");
+
+        // Create a new QueryManager (which will initialize a new PMF behind the scenes).
+        try (final var otherQm = new QueryManager()) {
+            // Request the global JDBI instance again and verify it differs from the original one.
+            // Because the PMF changed, a new instance must have been created.
+            final Jdbi otherJdbi = JdbiFactory.jdbi(otherQm);
+            assertThat(otherJdbi).isNotEqualTo(jdbi);
+
+            // Issue a test query to ensure the new JDBI instance is functional.
+            final Integer queryResult = otherJdbi.withHandle(handle ->
+                    handle.createQuery("SELECT 666").mapTo(Integer.class).one());
+            assertThat(queryResult).isEqualTo(666);
+        }
+    }
+
+    @Test
+    public void testLocalInstanceOutsideOfJdoTransaction() {
+        assertThatExceptionOfType(IllegalStateException.class)
+                .isThrownBy(() -> JdbiFactory.localJdbi(qm))
+                .withMessageContaining("Local JDBI instances must not be used outside of an active JDO transaction");
+    }
+
+    @Test
+    public void testLocalInstanceWithJdoTransaction() {
+        qm.runInTransaction(() -> {
+            // Create a new project.
+            final var project = new Project();
+            project.setName("acme-app");
+            project.setVersion("1.0.0");
+            qm.getPersistenceManager().makePersistent(project);
+
+            // Query for the created project, despite its creation not having been committed yet.
+            // Because the local JDBI instance uses the same connection as the QueryManager,
+            // it must be able to see the yet-uncommitted change.
+            final Optional<String> projectName = JdbiFactory.localJdbi(qm).withHandle(handle ->
+                    handle.createQuery("SELECT \"NAME\" FROM \"PROJECT\"").mapTo(String.class).findFirst());
+            assertThat(projectName).contains("acme-app");
+
+            // Ensure the connection is still usable after being returned from JDBI,
+            // by creating another record using the QueryManager.
+            final var component = new Component();
+            component.setProject(project);
+            component.setName("acme-lib");
+            component.setVersion("2.0.0");
+            qm.getPersistenceManager().makePersistent(component);
+        });
+    }
+
+}