SA

Author: marko-kriskovic

server/libs/automation/automation-knowledge-base/automation-knowledge-base-graphql/src/test/java/com/bytechef/automation/knowledgebase/web/graphql/KnowledgeBaseDocumentGraphQlControllerIntTest.java

@@ -95,8 +95,8 @@ void testGetKnowledgeBaseDocumentChunks() {
         KnowledgeBaseDocument mockDocument = createMockDocument(documentId, "Test Document");
 
         List<KnowledgeBaseDocumentChunk> mockChunks = List.of(
-            createMockChunk(1L, "Chunk 1"),
-            createMockChunk(2L, "Chunk 2"));
+            createMockChunk(1L),
+            createMockChunk(2L));
 
         when(knowledgeBaseDocumentService.getKnowledgeBaseDocument(documentId)).thenReturn(mockDocument);
         when(knowledgeBaseDocumentChunkFacade.getKnowledgeBaseDocumentChunksByDocumentId(documentId))
@@ -222,7 +222,7 @@ private KnowledgeBaseDocument createMockDocument(Long id, String name) {
         return document;
     }
 
-    private KnowledgeBaseDocumentChunk createMockChunk(Long id, String content) {
+    private KnowledgeBaseDocumentChunk createMockChunk(Long id) {
         KnowledgeBaseDocumentChunk chunk = new KnowledgeBaseDocumentChunk();
 
         chunk.setId(id);

server/libs/automation/automation-knowledge-base/automation-knowledge-base-service/src/main/java/com/bytechef/automation/knowledgebase/config/KnowledgeBaseVectorStoreMetadataUpdater.java

@@ -18,9 +18,11 @@
 
 import static com.bytechef.automation.knowledgebase.constant.KnowledgeBaseConstants.METADATA_TAG_NAMES;
 
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Pattern;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.jdbc.core.JdbcTemplate;
@@ -32,17 +34,24 @@
  *
  * @author Marko Kriskovic
  */
-public class KnowledgeBaseVectorStoreMetadataUpdater {
+public final class KnowledgeBaseVectorStoreMetadataUpdater {
 
     private static final Logger logger = LoggerFactory.getLogger(KnowledgeBaseVectorStoreMetadataUpdater.class);
+    private static final Pattern SAFE_TABLE_NAME_PATTERN = Pattern.compile("^[a-zA-Z_][a-zA-Z0-9_.]*$");
 
     private final JdbcTemplate pgVectorJdbcTemplate;
     private final ObjectMapper objectMapper;
     private final String fullTableName;
 
+    @SuppressFBWarnings("EI2")
     public KnowledgeBaseVectorStoreMetadataUpdater(
         JdbcTemplate pgVectorJdbcTemplate, ObjectMapper objectMapper, String fullTableName) {
 
+        if (!SAFE_TABLE_NAME_PATTERN.matcher(fullTableName)
+            .matches()) {
+            throw new IllegalArgumentException("Invalid table name: " + fullTableName);
+        }
+
         this.pgVectorJdbcTemplate = pgVectorJdbcTemplate;
         this.objectMapper = objectMapper;
         this.fullTableName = fullTableName;
@@ -56,6 +65,7 @@ public KnowledgeBaseVectorStoreMetadataUpdater(
      * @param vectorStoreId the vector store document ID
      * @param tagNames      the new tag names to set; an empty list removes all tag fields
      */
+    @SuppressFBWarnings("SQL_INJECTION_SPRING_JDBC")
     public void updateTagNames(String vectorStoreId, List<String> tagNames) {
         List<String> rows = pgVectorJdbcTemplate.queryForList(
             "SELECT metadata::text FROM " + fullTableName + " WHERE id = ?::uuid",

server/libs/automation/automation-knowledge-base/automation-knowledge-base-service/src/test/java/com/bytechef/automation/knowledgebase/config/KnowledgeBaseIntTestConfiguration.java

@@ -93,6 +93,11 @@ FileStorageServiceRegistry fileStorageServiceRegistry(FileStorageService fileSto
         return registry;
     }
 
+    @Bean
+    KnowledgeBaseVectorStoreMetadataUpdater knowledgeBaseVectorStoreMetadataUpdater() {
+        return mock(KnowledgeBaseVectorStoreMetadataUpdater.class);
+    }
+
     @Bean
     MessageBroker messageBroker() {
         return mock(MessageBroker.class);