feat: enhance error logging and add error_handler to queries

aperturedb/CommonLibrary.py

@@ -271,7 +271,8 @@ def execute_query(client: Connector, query: Commands,
                   blobs: Blobs = [],
                   success_statuses: list[int] = [0],
                   response_handler: Optional[Callable] = None, commands_per_query: int = 1, blobs_per_query: int = 0,
-                  strict_response_validation: bool = False, cmd_index=None) -> Tuple[int, CommandResponses, Blobs]:
+                  strict_response_validation: bool = False, cmd_index=None,
+                  error_handler: Optional[Callable] = None) -> Tuple[int, CommandResponses, Blobs]:
     """
     Execute a batch of queries, doing useful logging around it.
     Calls the response handler if provided.
@@ -288,6 +289,8 @@ def execute_query(client: Connector, query: Commands,
         commands_per_query (int, optional): The number of commands per query. Defaults to 1.
         blobs_per_query (int, optional): The number of blobs per query. Defaults to 0.
         strict_response_validation (bool, optional): Whether to strictly validate the response. Defaults to False.
+        cmd_index (int, optional): The index of the command or batch, passed to the response handler. Defaults to None.
+        error_handler (Callable, optional): Callback invoked when the query returns an unexpected status or fails. Expected signature is `error_handler(query, response, blobs)`. Defaults to None.
 
     Returns:
         int: The result code.
@@ -310,14 +313,21 @@ def execute_query(client: Connector, query: Commands,
                 map_response_to_handler(response_handler,
                                         query, blobs, r, b, commands_per_query, blobs_per_query,
                                         cmd_index)
-            except BaseException as e:
+            except Exception as e:
                 logger.exception(e)
                 if strict_response_validation:
-                    raise e
+                    raise
     else:
         # Transaction failed entirely.
+        num_commands = len(query) if isinstance(query, list) else 1
+        truncated_query = query[:2] if isinstance(
+            query, list) and num_commands > 2 else query
+        query_summary = (
+            f"{num_commands} commands (showing first 2: {truncated_query})"
+            if num_commands > 2 else str(query)
+        )
         logger.error(
-            f"Failed query = {query} with response = {censor_tokens(r)}")
+            f"Failed query = {query_summary} with response = {censor_tokens(r)}")
         result = 1
 
     statuses = {}
@@ -334,16 +344,25 @@ def execute_query(client: Connector, query: Commands,
 
     # last_query_ok means result status >= 0
     if result != 1:
-        warn_list = []
+        warnings_count = 0
         for status, results in statuses.items():
             if status not in success_statuses:
-                for wr in results:
-                    warn_list.append(wr)
-        if len(warn_list) != 0:
+                warnings_count += len(results)
+        if warnings_count != 0:
             logger.warning(
-                f"Partial errors:\r\n{json.dumps(query, default=str)}\r\n{json.dumps(censor_tokens(warn_list), default=str)}")
+                f"Encountered {warnings_count} partial errors. "
+                "Use error_handler or inspect response for details."
+            )
             result = 2
 
+    if result != 0 and error_handler is not None:
+        try:
+            error_handler(query, r, blobs)
+        except Exception as e:
+            logger.exception(e)
+            if strict_response_validation:
+                raise
+
     return result, r, b
 
 

aperturedb/ParallelQuery.py

@@ -158,9 +158,12 @@ def process_responses(requests, input_blobs, responses, output_blobs):
         worker_stats = {}
         if not self.dry_run:
             response_handler = None
+            error_handler = None
             strict_response_validation = False
             if hasattr(self.generator, "response_handler") and callable(self.generator.response_handler):
                 response_handler = self.generator.response_handler
+            if hasattr(self.generator, "error_handler") and callable(self.generator.error_handler):
+                error_handler = self.generator.error_handler
             if hasattr(self.generator, "strict_response_validation") and isinstance(self.generator.strict_response_validation, bool):
                 strict_response_validation = self.generator.strict_response_validation
 
@@ -186,7 +189,8 @@ def response_handler(query, qblobs, resp, rblobs, qindex): return indexless_hand
                 self.commands_per_query,
                 self.blobs_per_query,
                 strict_response_validation=strict_response_validation,
-                cmd_index=batch_start)
+                cmd_index=batch_start,
+                error_handler=error_handler)
             if result == 0:
                 query_time = client.get_last_query_time()
                 worker_stats["succeeded_commands"] = len(q)

aperturedb/ParallelQuerySet.py

@@ -46,23 +46,32 @@ def gen_execute_batch_sets(base_executor):
     # if blob_set is a list of lists, each inner list will be given to the inner
     #  execution
     #
-    def execute_batch_sets(client, query_set, blob_set, success_statuses: list[int] = [0],
-                           response_handler: Optional[Callable] = None, commands_per_query: list[int] = -1,
-                           blobs_per_query: list[int] = -1,
-                           strict_response_validation: bool = False, cmd_index: int = None):
+    def execute_batch_sets(client, query_set, blob_set, success_statuses: Optional[list[int]] = None,
+                           response_handler: Optional[Callable] = None, commands_per_query: Optional[list[int]] = None,
+                           blobs_per_query: Optional[list[int]] = None,
+                           strict_response_validation: bool = False, cmd_index: Optional[int] = None,
+                           error_handler: Optional[Callable] = None):
 
-        logger.info("Execute Batch Sets = Batch Size {0}  Comands Per Query {1} Blobs Per Query {2}".format(
-            len(query_set), commands_per_query, blobs_per_query))
+        if success_statuses is None:
+            success_statuses = ParallelQuery.success_statuses
 
         batch_size = len(query_set)
 
         # test query set
         first_element = query_set[0]
         if not isinstance(first_element, list):
-            logger.error("First Element not a list: {first_element}")
+            logger.error(f"First Element not a list: {first_element}")
             raise Exception("Query set must be a list of lists")
         set_total = len(first_element)
 
+        if commands_per_query is None:
+            commands_per_query = [1] * set_total
+        if blobs_per_query is None:
+            blobs_per_query = [0] * set_total
+
+        logger.info("Execute Batch Sets = Batch Size {0}  Commands Per Query {1} Blobs Per Query {2}".format(
+            batch_size, commands_per_query, blobs_per_query))
+
         # Check if blobs are a simple array or nested array of blobs
         per_set_blobs = isinstance(blob_set, list) and len(
             blob_set) > 0 and isinstance(blob_set[0], list)
@@ -273,7 +282,7 @@ def constraint_filter(single_line, single_results):
 
                 query_filter = constraint_filter
 
-            local_success_statuses = [0, 2]
+            local_success_statuses = success_statuses
 
             # queries are by row first, so we run query_filter on each query
             # we pass the entire row's data, then we retrieve all of the stored results for that row
@@ -303,7 +312,8 @@ def constraint_filter(single_line, single_results):
                                                                   commands_per_query[i],
                                                                   blobs_per_query[i],
                                                                   strict_response_validation=strict_response_validation,
-                                                                  cmd_index=cmd_index)
+                                                                  cmd_index=cmd_index,
+                                                                  error_handler=error_handler)
                 if response_handler != None and client.last_query_ok():
                     def map_to_set(query, query_blobs, resp, resp_blobs):
                         response_handler(

aperturedb/VideoDataCSV.py

@@ -136,9 +136,8 @@ def load_video(self, filename):
                 logger.exception(f"Video Error: {filename}")
 
         try:
-            fd = open(filename, "rb")
-            buff = fd.read()
-            fd.close()
+            with open(filename, "rb") as fd:
+                buff = fd.read()
             return True, buff
         except Exception:
             logger.exception(f"Video Error: {filename}")

test/test_ResponseHandler.py

@@ -488,3 +488,104 @@ def mock_query(self, request, blobs):
         querier.query(generator, numthreads=1, batchsize=2, stats=False)
 
         assert len(changed_ids) == 2
+
+    def test_error_handler_total_failure(self, db, monkeypatch):
+        from aperturedb.QueryGenerator import QueryGenerator
+        from aperturedb.ParallelQuery import ParallelQuery
+        from aperturedb.Connector import Connector
+
+        class QGErrorHandlerFailure(QueryGenerator):
+            def __init__(self):
+                self.error_handler_called = False
+
+            def __len__(self):
+                return 1
+
+            def getitem(self, idx):
+                return [{"FindImage": {}}], []
+
+            def error_handler(self, q, r, b):
+                self.error_handler_called = True
+
+        def mock_query(self, request, blobs):
+            self.response = {"status": -1, "info": "Transaction failed"}
+            self.blobs = []
+            return self.response, []
+
+        monkeypatch.setattr(Connector, "query", mock_query)
+        monkeypatch.setattr(Connector, "last_query_ok", lambda self: False)
+        monkeypatch.setattr(Connector, "clone", lambda self: self)
+
+        generator = QGErrorHandlerFailure()
+        querier = ParallelQuery(db)
+        querier.query(generator, numthreads=1, batchsize=1, stats=False)
+
+        assert generator.error_handler_called
+
+    def test_error_handler_partial_error(self, db, monkeypatch):
+        from aperturedb.QueryGenerator import QueryGenerator
+        from aperturedb.ParallelQuery import ParallelQuery
+        from aperturedb.Connector import Connector
+
+        class QGErrorHandlerPartial(QueryGenerator):
+            def __init__(self):
+                self.error_handler_called = False
+
+            def __len__(self):
+                return 1
+
+            def getitem(self, idx):
+                return [{"FindImage": {}}, {"AddImage": {}}], []
+
+            def error_handler(self, q, r, b):
+                self.error_handler_called = True
+
+        def mock_query(self, request, blobs):
+            self.response = [{"FindImage": {"status": 0}},
+                             {"AddImage": {"status": 1}}]
+            self.blobs = []
+            return self.response, []
+
+        monkeypatch.setattr(Connector, "query", mock_query)
+        monkeypatch.setattr(Connector, "last_query_ok", lambda self: True)
+        monkeypatch.setattr(Connector, "clone", lambda self: self)
+
+        generator = QGErrorHandlerPartial()
+        querier = ParallelQuery(db)
+        querier.query(generator, numthreads=1, batchsize=1, stats=False)
+
+        assert generator.error_handler_called
+
+    def test_error_handler_success(self, db, monkeypatch):
+        from aperturedb.QueryGenerator import QueryGenerator
+        from aperturedb.ParallelQuery import ParallelQuery
+        from aperturedb.Connector import Connector
+
+        class QGErrorHandlerSuccess(QueryGenerator):
+            def __init__(self):
+                self.error_handler_called = False
+
+            def __len__(self):
+                return 1
+
+            def getitem(self, idx):
+                return [{"FindImage": {}}, {"AddImage": {}}], []
+
+            def error_handler(self, q, r, b):
+                self.error_handler_called = True
+
+        def mock_query(self, request, blobs):
+            self.response = [{"FindImage": {"status": 0}},
+                             {"AddImage": {"status": 0}}]
+            self.blobs = []
+            return self.response, []
+
+        monkeypatch.setattr(Connector, "query", mock_query)
+        monkeypatch.setattr(Connector, "last_query_ok", lambda self: True)
+        monkeypatch.setattr(Connector, "clone", lambda self: self)
+
+        generator = QGErrorHandlerSuccess()
+        querier = ParallelQuery(db)
+        querier.query(generator, numthreads=1, batchsize=1, stats=False)
+
+        assert not generator.error_handler_called