Merge pull request #271 from cloudera/main

Release 1.25.0

Author: ewilliams-cloudera

backend/src/main/java/com/cloudera/cai/rag/configuration/DatabaseOperations.java

@@ -0,0 +1,71 @@
+/*
+ * CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
+ * (C) Cloudera, Inc. 2024
+ * All rights reserved.
+ *
+ * Applicable Open Source License: Apache 2.0
+ *
+ * NOTE: Cloudera open source products are modular software products
+ * made up of hundreds of individual components, each of which was
+ * individually copyrighted.  Each Cloudera open source product is a
+ * collective work under U.S. Copyright Law. Your license to use the
+ * collective work is as provided in your written agreement with
+ * Cloudera.  Used apart from the collective work, this file is
+ * licensed for your use pursuant to the open source license
+ * identified above.
+ *
+ * This code is provided to you pursuant a written agreement with
+ * (i) Cloudera, Inc. or (ii) a third-party authorized to distribute
+ * this code. If you do not have a written agreement with Cloudera nor
+ * with an authorized and properly licensed third party, you do not
+ * have any rights to access nor to use this code.
+ *
+ * Absent a written agreement with Cloudera, Inc. ("Cloudera") to the
+ * contrary, A) CLOUDERA PROVIDES THIS CODE TO YOU WITHOUT WARRANTIES OF ANY
+ * KIND; (B) CLOUDERA DISCLAIMS ANY AND ALL EXPRESS AND IMPLIED
+ * WARRANTIES WITH RESPECT TO THIS CODE, INCLUDING BUT NOT LIMITED TO
+ * IMPLIED WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE; (C) CLOUDERA IS NOT LIABLE TO YOU,
+ * AND WILL NOT DEFEND, INDEMNIFY, NOR HOLD YOU HARMLESS FOR ANY CLAIMS
+ * ARISING FROM OR RELATED TO THE CODE; AND (D)WITH RESPECT TO YOUR EXERCISE
+ * OF ANY RIGHTS GRANTED TO YOU FOR THE CODE, CLOUDERA IS NOT LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, PUNITIVE OR
+ * CONSEQUENTIAL DAMAGES INCLUDING, BUT NOT LIMITED TO, DAMAGES
+ * RELATED TO LOST REVENUE, LOST PROFITS, LOSS OF INCOME, LOSS OF
+ * BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
+ * DATA.
+ ******************************************************************************/
+
+package com.cloudera.cai.rag.configuration;
+
+import org.jdbi.v3.core.HandleCallback;
+import org.jdbi.v3.core.HandleConsumer;
+import org.jdbi.v3.core.Jdbi;
+
+/**
+ * Database operations implementation that delegates to a Jdbi instance. This is the production
+ * implementation used throughout the application.
+ */
+public class DatabaseOperations {
+  private final Jdbi jdbi;
+
+  public DatabaseOperations(Jdbi jdbi) {
+    this.jdbi = jdbi;
+  }
+
+  public <X extends Exception> void useHandle(HandleConsumer<X> handleConsumer) throws X {
+    jdbi.useHandle(handleConsumer);
+  }
+
+  public <X extends Exception> void useTransaction(HandleConsumer<X> handleConsumer) throws X {
+    jdbi.useTransaction(handleConsumer);
+  }
+
+  public <T, X extends Exception> T inTransaction(HandleCallback<T, X> handleCallback) throws X {
+    return jdbi.inTransaction(handleCallback);
+  }
+
+  public <T, X extends Exception> T withHandle(HandleCallback<T, X> handleCallback) throws X {
+    return jdbi.withHandle(handleCallback);
+  }
+}

backend/src/main/java/com/cloudera/cai/rag/configuration/JdbiConfiguration.java

@@ -1,4 +1,4 @@
-/*******************************************************************************
+/*
  * CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
  * (C) Cloudera, Inc. 2024
  * All rights reserved.
@@ -45,6 +45,8 @@
 import java.sql.SQLException;
 import javax.sql.DataSource;
 import lombok.extern.slf4j.Slf4j;
+import org.jdbi.v3.core.HandleCallback;
+import org.jdbi.v3.core.HandleConsumer;
 import org.jdbi.v3.core.Jdbi;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -56,8 +58,8 @@ public class JdbiConfiguration {
   private static final Object LOCK = new Object();
 
   @Bean
-  public Jdbi jdbi() {
-    return createJdbi();
+  public DatabaseOperations databaseOperations() {
+    return new DatabaseOperations(createJdbi());
   }
 
   private static Jdbi createJdbi() {
@@ -108,7 +110,51 @@ private static DataSource dataSource(DatabaseConfig databaseConfig) throws SQLEx
   }
 
   // nullables below here
-  public static Jdbi createNull() {
-    return new JdbiConfiguration().jdbi();
+  public static DatabaseOperations createNull(RuntimeException... exceptions) {
+    return new DatabaseOperationsStub(createJdbi(), exceptions);
+  }
+
+  /**
+   * Test implementation of DatabaseOperations that can inject failures for testing. This allows us
+   * to test error scenarios without using Mockito.
+   */
+  private static class DatabaseOperationsStub extends DatabaseOperations {
+    private final RuntimeException[] exceptions;
+    private int exceptionIndex = 0;
+
+    private DatabaseOperationsStub(Jdbi jdbi, RuntimeException[] exceptions) {
+      super(jdbi);
+      this.exceptions = exceptions;
+    }
+
+    private void checkForException() {
+      if (exceptionIndex < exceptions.length) {
+        throw exceptions[exceptionIndex++];
+      }
+    }
+
+    @Override
+    public <X extends Exception> void useHandle(HandleConsumer<X> handleConsumer) throws X {
+      checkForException();
+      super.useHandle(handleConsumer);
+    }
+
+    @Override
+    public <X extends Exception> void useTransaction(HandleConsumer<X> handleConsumer) throws X {
+      checkForException();
+      super.useTransaction(handleConsumer);
+    }
+
+    @Override
+    public <T, X extends Exception> T inTransaction(HandleCallback<T, X> handleCallback) throws X {
+      checkForException();
+      return super.inTransaction(handleCallback);
+    }
+
+    @Override
+    public <T, X extends Exception> T withHandle(HandleCallback<T, X> handleCallback) throws X {
+      checkForException();
+      return super.withHandle(handleCallback);
+    }
   }
 }

backend/src/main/java/com/cloudera/cai/rag/datasources/DeleteDataSourceReconciler.java

@@ -38,37 +38,37 @@
 
 package com.cloudera.cai.rag.datasources;
 
+import com.cloudera.cai.rag.configuration.DatabaseOperations;
 import com.cloudera.cai.rag.configuration.JdbiConfiguration;
 import com.cloudera.cai.rag.external.RagBackendClient;
 import com.cloudera.cai.util.Tracker;
 import com.cloudera.cai.util.reconcilers.*;
 import io.opentelemetry.api.OpenTelemetry;
 import java.util.Set;
 import lombok.extern.slf4j.Slf4j;
-import org.jdbi.v3.core.Jdbi;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.stereotype.Component;
 
 @Component
 @Slf4j
 public class DeleteDataSourceReconciler extends BaseReconciler<Long> {
-  private final Jdbi jdbi;
+  private final DatabaseOperations databaseOperations;
   private final RagBackendClient ragBackendClient;
 
   public DeleteDataSourceReconciler(
-      Jdbi jdbi,
+      DatabaseOperations databaseOperations,
       RagBackendClient ragBackendClient,
       @Qualifier("reconcilerConfig") ReconcilerConfig reconcilerConfig,
       OpenTelemetry openTelemetry) {
     super(reconcilerConfig, openTelemetry);
-    this.jdbi = jdbi;
+    this.databaseOperations = databaseOperations;
     this.ragBackendClient = ragBackendClient;
   }
 
   @Override
   public void resync() {
     log.debug("Checking for data sources to delete");
-    jdbi.useHandle(
+    databaseOperations.useHandle(
         handle ->
             handle
                 .createQuery("SELECT id FROM rag_data_source WHERE deleted IS NOT NULL")
@@ -83,7 +83,7 @@ public ReconcileResult reconcile(Set<Long> dataSourceIds) {
       ragBackendClient.deleteDataSource(dataSourceId);
       log.info(
           "deleting data source and documents from the database. data source id: {}", dataSourceId);
-      jdbi.useTransaction(
+      databaseOperations.useTransaction(
           handle -> {
             handle.execute("DELETE FROM RAG_DATA_SOURCE WHERE ID = ?", dataSourceId);
             handle.execute(

backend/src/main/java/com/cloudera/cai/rag/datasources/RagDataSourceRepository.java

@@ -39,28 +39,28 @@
 package com.cloudera.cai.rag.datasources;
 
 import com.cloudera.cai.rag.Types.RagDataSource;
+import com.cloudera.cai.rag.configuration.DatabaseOperations;
 import com.cloudera.cai.rag.configuration.JdbiConfiguration;
 import com.cloudera.cai.util.exceptions.NotFound;
 import java.time.Instant;
 import java.util.List;
 import lombok.extern.slf4j.Slf4j;
 import org.jdbi.v3.core.Handle;
-import org.jdbi.v3.core.Jdbi;
 import org.jdbi.v3.core.mapper.reflect.ConstructorMapper;
 import org.jdbi.v3.core.statement.Query;
 import org.springframework.stereotype.Component;
 
 @Component
 @Slf4j
 public class RagDataSourceRepository {
-  private final Jdbi jdbi;
+  private final DatabaseOperations databaseOperations;
 
-  public RagDataSourceRepository(Jdbi jdbi) {
-    this.jdbi = jdbi;
+  public RagDataSourceRepository(DatabaseOperations databaseOperations) {
+    this.databaseOperations = databaseOperations;
   }
 
   public Long createRagDataSource(RagDataSource input) {
-    return jdbi.inTransaction(handle -> createRagDataSource(handle, input));
+    return databaseOperations.inTransaction(handle -> createRagDataSource(handle, input));
   }
 
   public Long createRagDataSource(Handle handle, RagDataSource input) {
@@ -91,7 +91,7 @@ private static RagDataSource cleanInputs(RagDataSource input) {
 
   public void updateRagDataSource(RagDataSource input) {
     RagDataSource cleanedInputs = cleanInputs(input);
-    jdbi.useTransaction(
+    databaseOperations.useTransaction(
         handle -> {
           var sql =
               """
@@ -121,7 +121,7 @@ public void updateRagDataSource(RagDataSource input) {
   }
 
   public RagDataSource getRagDataSourceById(Long id) {
-    return jdbi.withHandle(
+    return databaseOperations.withHandle(
         handle -> {
           var sql =
               """
@@ -150,7 +150,7 @@ public RagDataSource getRagDataSourceById(Long id) {
 
   public List<RagDataSource> getRagDataSources() {
     log.debug("Getting all RagDataSources");
-    return jdbi.withHandle(
+    return databaseOperations.withHandle(
         handle -> {
           var sql =
               """
@@ -174,7 +174,7 @@ public List<RagDataSource> getRagDataSources() {
   }
 
   public void deleteDataSource(Long id) {
-    jdbi.useTransaction(handle -> deleteDataSource(handle, id));
+    databaseOperations.useTransaction(handle -> deleteDataSource(handle, id));
   }
 
   public void deleteDataSource(Handle handle, Long id) {
@@ -184,7 +184,7 @@ public void deleteDataSource(Handle handle, Long id) {
   }
 
   public int getNumberOfDataSources() {
-    return jdbi.withHandle(
+    return databaseOperations.withHandle(
         handle -> {
           try (var query =
               handle.createQuery(
@@ -198,6 +198,6 @@ public int getNumberOfDataSources() {
 
   public static RagDataSourceRepository createNull() {
     // the db configuration will use in-memory db based on env vars.
-    return new RagDataSourceRepository(new JdbiConfiguration().jdbi());
+    return new RagDataSourceRepository(JdbiConfiguration.createNull());
   }
 }

backend/src/main/java/com/cloudera/cai/rag/files/RagFileDeleteReconciler.java

@@ -1,6 +1,7 @@
 package com.cloudera.cai.rag.files;
 
 import com.cloudera.cai.rag.Types;
+import com.cloudera.cai.rag.configuration.DatabaseOperations;
 import com.cloudera.cai.rag.configuration.JdbiConfiguration;
 import com.cloudera.cai.rag.external.RagBackendClient;
 import com.cloudera.cai.util.exceptions.NotFound;
@@ -10,31 +11,30 @@
 import io.opentelemetry.api.OpenTelemetry;
 import java.util.Set;
 import lombok.extern.slf4j.Slf4j;
-import org.jdbi.v3.core.Jdbi;
 import org.jdbi.v3.core.mapper.reflect.ConstructorMapper;
 import org.jdbi.v3.core.statement.Query;
 import org.springframework.stereotype.Component;
 
 @Slf4j
 @Component
 public class RagFileDeleteReconciler extends BaseReconciler<Types.RagDocument> {
-  private final Jdbi jdbi;
+  private final DatabaseOperations databaseOperations;
   private final RagBackendClient ragBackendClient;
 
   public RagFileDeleteReconciler(
       ReconcilerConfig reconcilerConfig,
       OpenTelemetry openTelemetry,
-      Jdbi jdbi,
+      DatabaseOperations databaseOperations,
       RagBackendClient ragBackendClient) {
     super(reconcilerConfig, openTelemetry);
-    this.jdbi = jdbi;
+    this.databaseOperations = databaseOperations;
     this.ragBackendClient = ragBackendClient;
   }
 
   @Override
   public void resync() throws Exception {
     log.debug("checking for RAG documents to be deleted");
-    jdbi.useHandle(
+    databaseOperations.useHandle(
         handle -> {
           handle.registerRowMapper(ConstructorMapper.factory(Types.RagDocument.class));
           Query query =
@@ -54,7 +54,7 @@ public ReconcileResult reconcile(Set<Types.RagDocument> documents) throws Except
       } catch (NotFound e) {
         log.debug("got a not found exception from the rag backend: {}", e.getMessage());
       }
-      jdbi.useHandle(
+      databaseOperations.useHandle(
           handle ->
               handle.execute("DELETE from rag_data_source_document WHERE id = ?", document.id()));
     }