diff --git a/src/python/txtai/database/duckdb.py b/src/python/txtai/database/duckdb.py index b2a51cee..88cf6310 100644 --- a/src/python/txtai/database/duckdb.py +++ b/src/python/txtai/database/duckdb.py @@ -18,6 +18,7 @@ except ImportError: from .embedded import Embedded from .schema import Statement +from .sql import SQLError class DuckDB(Embedded): @@ -29,6 +30,11 @@ class DuckDB(Embedded): DELETE_DOCUMENT = "DELETE FROM documents WHERE id = ?" DELETE_OBJECT = "DELETE FROM objects WHERE id = ?" + # Added to errors that abort the transaction + ABORTED = ( + "This error aborted the DuckDB transaction. Unsaved changes are lost and later operations fail until the index is loaded or built again." + ) + def __init__(self, config): super().__init__(config) @@ -36,8 +42,15 @@ class DuckDB(Embedded): raise ImportError('DuckDB is not available - install "database" extra to enable') def execute(self, function, *args): - # Call parent method with DuckDB compatible arguments - return super().execute(function, *self.formatargs(args)) + try: + # Call parent method with DuckDB compatible arguments + return super().execute(function, *self.formatargs(args)) + except SQLError as e: + # Runtime errors abort the transaction. DuckDB has no savepoints, so it can't be recovered without a rollback. + if self.aborted(): + raise SQLError(f"{e}\n\n{DuckDB.ABORTED}") from None + + raise def insertdocument(self, uid, data, tags, entry): # Delete existing document @@ -123,6 +136,20 @@ class DuckDB(Embedded): return connection + def aborted(self): + """ + Checks if the current transaction is aborted. + + Returns: + True if the current transaction is aborted + """ + + try: + self.connection.execute("SELECT 1") + return False + except duckdb.TransactionException: + return True + def formatargs(self, args): """ DuckDB doesn't support named parameters. This method replaces named parameters with question marks diff --git a/test/python/testdatabase/testduckdb.py b/test/python/testdatabase/testduckdb.py index 677c2eb9..b4e45e0b 100644 --- a/test/python/testdatabase/testduckdb.py +++ b/test/python/testdatabase/testduckdb.py @@ -5,6 +5,7 @@ DuckDB module tests import os import unittest +from txtai.database import DuckDB, SQLError from txtai.embeddings import Embeddings from .testrdbms import Common @@ -45,6 +46,28 @@ class TestDuckDB(Common.TestRDBMS): if cls.embeddings: cls.embeddings.close() + def testAborted(self): + """ + Test SQL errors that abort the transaction + """ + + embeddings = Embeddings(keyword=True, content=self.backend) + embeddings.index([(0, "apple pie", None), (1, "banana bread", None)]) + + # Catalog errors don't abort the transaction + with self.assertRaises(SQLError) as context: + embeddings.search("select id from txtai where unknownfn(text)") + + self.assertNotIn(DuckDB.ABORTED, str(context.exception)) + + # Runtime errors abort the transaction + with self.assertRaises(SQLError) as context: + embeddings.search("select id from txtai where text = :x", parameters={"x": 1}) + + self.assertIn("This error aborted the DuckDB transaction.", str(context.exception)) + + embeddings.close() + @unittest.skipIf(os.name == "nt", "testArchive skipped on Windows") def testArchive(self): """