diff --git a/aperturedb/CommonLibrary.py b/aperturedb/CommonLibrary.py index e147a0a3..ab91794a 100644 --- a/aperturedb/CommonLibrary.py +++ b/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 diff --git a/aperturedb/ParallelQuery.py b/aperturedb/ParallelQuery.py index 498aed7e..e2237420 100644 --- a/aperturedb/ParallelQuery.py +++ 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) diff --git a/aperturedb/ParallelQuerySet.py b/aperturedb/ParallelQuerySet.py index 119379b7..9a6e885c 100644 --- a/aperturedb/ParallelQuerySet.py +++ b/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( diff --git a/aperturedb/VideoDataCSV.py b/aperturedb/VideoDataCSV.py index 39e4ee97..ef3135be 100644 --- a/aperturedb/VideoDataCSV.py +++ b/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}") diff --git a/test/test_ResponseHandler.py b/test/test_ResponseHandler.py index 91b6bb9d..3d01368a 100644 --- a/test/test_ResponseHandler.py +++ b/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