From 3d83f88fc35d312c89943131736310e085e6be30 Mon Sep 17 00:00:00 2001
From: Ralph Gasser <rg@pontius.ch>
Date: Thu, 11 Jul 2024 16:09:09 +0200
Subject: [PATCH] Adds remaining base operations to
 AbstractRetrievableWriterTest.

Signed-off-by: Ralph Gasser <rg@pontius.ch>
---
 .../AbstractRetrievableWriterTest.kt          | 54 +++++++++++++++++++
 .../pgvector/retrievable/RetrievableWriter.kt |  2 +-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/vitrivr-engine-core/src/testFixtures/kotlin/org/vitrivr/engine/core/database/retrievable/AbstractRetrievableWriterTest.kt b/vitrivr-engine-core/src/testFixtures/kotlin/org/vitrivr/engine/core/database/retrievable/AbstractRetrievableWriterTest.kt
index ecc256d1..4701ceeb 100644
--- a/vitrivr-engine-core/src/testFixtures/kotlin/org/vitrivr/engine/core/database/retrievable/AbstractRetrievableWriterTest.kt
+++ b/vitrivr-engine-core/src/testFixtures/kotlin/org/vitrivr/engine/core/database/retrievable/AbstractRetrievableWriterTest.kt
@@ -81,6 +81,60 @@ abstract class AbstractRetrievableWriterTest(schemaPath: String) : AbstractDatab
         }
     }
 
+    /**
+     * Tests if the [RetrievableWriter.add] works as expected.
+     */
+    @Test
+    fun testDeleteAll() {
+        val writer = this.testConnection.getRetrievableWriter()
+        val reader = this.testConnection.getRetrievableReader()
+        val size = Random().nextInt(500, 5000)
+
+        /* Create and add retrievable. */
+        val ids = (0 until size).map { UUID.randomUUID() }
+        val ingested = ids.map { Ingested(it, "INGESTED:TEST", false) }
+
+        /* Execute actions. */
+        Assertions.assertTrue(writer.addAll(ingested))
+
+        /* Check if retrievable can be read. */
+        Assertions.assertEquals(ids.size.toLong(), reader.count())
+        reader.getAll(ids).forEachIndexed() { i, it ->
+            Assertions.assertEquals("INGESTED:TEST", it.type)
+        }
+
+        /* Execute actions. */
+        Assertions.assertTrue(writer.deleteAll(ingested))
+        Assertions.assertEquals(0L, reader.count())
+    }
+
+
+    /**
+     * Tests if the [RetrievableWriter.add] works as expected.
+     */
+    @Test
+    fun testUpdate() {
+        val writer = this.testConnection.getRetrievableWriter()
+        val reader = this.testConnection.getRetrievableReader()
+        val size = Random().nextInt(500, 5000)
+
+        /* Create and add retrievable. */
+        val ids = (0 until size).map { UUID.randomUUID() }
+        val ingested = ids.map { Ingested(it, "INGESTED:TEST", false) }
+        val update = ingested[Random().nextInt(0, ingested.size)]
+
+        /* Execute actions. */
+        Assertions.assertTrue(writer.addAll(ingested))
+
+        /* Check if retrievable can be read. */
+        Assertions.assertEquals(ids.size.toLong(), reader.count())
+        Assertions.assertTrue(writer.update(update.copy(type = "INGESTED:TEST2")))
+
+        /* Execute actions. */
+        Assertions.assertEquals(ids.size.toLong(), reader.count())
+        Assertions.assertEquals("INGESTED:TEST2", reader[update.id]?.type)
+    }
+
     /**
      * Cleans up the database after each test.
      */
diff --git a/vitrivr-engine-module-pgvector/src/main/kotlin/org/vitrivr/engine/database/pgvector/retrievable/RetrievableWriter.kt b/vitrivr-engine-module-pgvector/src/main/kotlin/org/vitrivr/engine/database/pgvector/retrievable/RetrievableWriter.kt
index 33c7f8cf..5de14b54 100644
--- a/vitrivr-engine-module-pgvector/src/main/kotlin/org/vitrivr/engine/database/pgvector/retrievable/RetrievableWriter.kt
+++ b/vitrivr-engine-module-pgvector/src/main/kotlin/org/vitrivr/engine/database/pgvector/retrievable/RetrievableWriter.kt
@@ -63,7 +63,7 @@ internal class RetrievableWriter(override val connection: PgVectorConnection): R
             this.connection.jdbc.prepareStatement("UPDATE $RETRIEVABLE_ENTITY_NAME SET $RETRIEVABLE_TYPE_COLUMN_NAME = ? WHERE $RETRIEVABLE_ID_COLUMN_NAME = ?").use { stmt ->
                 stmt.setString(1, item.type)
                 stmt.setObject(2, item.id)
-                return stmt.execute()
+                return stmt.executeUpdate() == 1
             }
         } catch (e: SQLException) {
             LOGGER.error(e) { "Failed to update retrievable ${item.id} due to SQL error." }