From 31e2e7223ff82c58d184e5a872b9fb2c9ad84ce5 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Thu, 21 May 2026 11:12:50 -0700 Subject: [PATCH 01/19] feat: make dev container for developers (#709) --- .devcontainer/Dockerfile | 6 ++++++ .devcontainer/devcontainer.json | 1 + .devcontainer/docker-compose.yml | 3 ++- aperturedb/transformers/clip.py | 4 ++-- docker/devcontainer/Dockerfile | 5 ----- docker/notebook/Dockerfile.cpu | 2 +- 6 files changed, 12 insertions(+), 9 deletions(-) create mode 100644 .devcontainer/Dockerfile delete mode 100644 docker/devcontainer/Dockerfile diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 00000000..e532a8e8 --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,6 @@ +# Pull base image. +FROM aperturedata/aperturedb-notebook:dependencies + +RUN python -m pip install --no-cache-dir \ + awscli==1.45.11 \ + openai-clip diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index f0ac70d0..9bdc7cef 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -9,6 +9,7 @@ "vscode": { "extensions": [ "ms-python.python", + "ms-python.vscode-pylance", "ms-python.pylint", "ms-toolsai.jupyter" ] diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml index a4a1b5e9..e66875b9 100644 --- a/.devcontainer/docker-compose.yml +++ b/.devcontainer/docker-compose.yml @@ -74,7 +74,8 @@ services: - ./aperturedb/certificate:/etc/nginx/certificate devcontainer: build: - context: ../docker/devcontainer + context: . + dockerfile: Dockerfile environment: DB_HOST: lenz DB_PORT: 55551 diff --git a/aperturedb/transformers/clip.py b/aperturedb/transformers/clip.py index 73e3b780..e1b85124 100644 --- a/aperturedb/transformers/clip.py +++ b/aperturedb/transformers/clip.py @@ -6,9 +6,9 @@ logger = logging.getLogger(__name__) error_message = """ -CLIP transformer requires git+https://github.com/openai/CLIP.git and torch +CLIP transformer requires openai-clip and torch Install with: pip install aperturedb[complete], followed by explicit install of CLIP. -Can be done with : "pip install git+https://github.com/openai/CLIP.git" in the same +Can be done with : "pip install openai-clip" in the same venv as aperturedb. """ diff --git a/docker/devcontainer/Dockerfile b/docker/devcontainer/Dockerfile deleted file mode 100644 index 039b1e91..00000000 --- a/docker/devcontainer/Dockerfile +++ /dev/null @@ -1,5 +0,0 @@ -# Pull base image. -FROM aperturedata/aperturedb-notebook:dependencies - -RUN pip install awscli -RUN pip install git+https://github.com/openai/CLIP.git diff --git a/docker/notebook/Dockerfile.cpu b/docker/notebook/Dockerfile.cpu index 60c3ac3e..c4b4b880 100644 --- a/docker/notebook/Dockerfile.cpu +++ b/docker/notebook/Dockerfile.cpu @@ -27,7 +27,7 @@ RUN pip install torch torchvision --index-url https://download.pytorch.org/whl/c RUN pip install facenet-pytorch --no-deps # Install CLIP (for running transformers) -RUN pip install git+https://github.com/openai/CLIP.git +RUN pip install openai-clip RUN apt update && apt install -y curl && apt clean From e672d7b8f965956a7b921632db3862860ccab2a9 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Thu, 21 May 2026 15:23:00 -0700 Subject: [PATCH 02/19] feat: add support for Tensorflow dataset (#729) --- aperturedb/PyTorchDataset.py | 18 +++- aperturedb/TensorFlowDataset.py | 173 ++++++++++++++++++++++++++++++++ test/test_tf_connector.py | 139 +++++++++++++++++++++++++ test/test_torch_connector.py | 2 +- 4 files changed, 330 insertions(+), 2 deletions(-) create mode 100644 aperturedb/TensorFlowDataset.py create mode 100644 test/test_tf_connector.py diff --git a/aperturedb/PyTorchDataset.py b/aperturedb/PyTorchDataset.py index f1d23670..ed29fe65 100644 --- a/aperturedb/PyTorchDataset.py +++ b/aperturedb/PyTorchDataset.py @@ -46,6 +46,13 @@ def __init__(self, client: Connector, query, label_prop=None, batch_size=1): if not "results" in self.query[self.find_image_idx]["FindImage"]: self.query[self.find_image_idx]["FindImage"]["results"] = {} + if self.label_prop is not None: + results = self.query[self.find_image_idx]["FindImage"]["results"] + if "list" not in results: + results["list"] = [] + if self.label_prop not in results["list"]: + results["list"].append(self.label_prop) + self.query[self.find_image_idx]["FindImage"]["batch"] = {} self.query[self.find_image_idx]["FindImage"]["blobs"] = True @@ -131,7 +138,16 @@ def get_batch(self, index): if self.label_prop: entities = r[self.find_image_idx]["FindImage"]["entities"] - self.batch_labels = [l[self.label_prop] for l in entities] + try: + self.batch_labels = [l[self.label_prop] for l in entities] + except KeyError: + prop = self.label_prop + msg = ( + f"Property '{prop}' not found in some entities. " + "Ensure all entities have this property." + ) + logger.error(msg) + raise else: self.batch_labels = ["none" for l in range(len(b))] except: diff --git a/aperturedb/TensorFlowDataset.py b/aperturedb/TensorFlowDataset.py new file mode 100644 index 00000000..12584457 --- /dev/null +++ b/aperturedb/TensorFlowDataset.py @@ -0,0 +1,173 @@ +import math +import numpy as np +import cv2 +import logging + +from aperturedb.CommonLibrary import execute_query +from aperturedb.Connector import Connector + +logger = logging.getLogger(__name__) + + +class ApertureDBTensorFlowDataset: + """ + This class implements a TensorFlow Dataset for ApertureDB. + It is used to load images from ApertureDB into a TensorFlow model. + It can be initialized with a query that will be used to retrieve + the images from ApertureDB. + """ + + def __init__(self, client: Connector, query, label_prop=None, batch_size=1): + + self.client = client.clone() + self.query = query + self.find_image_idx = None + self.total_elements = 0 + self.batch_size = batch_size + self.batch_images = [] + self.batch_start = 0 + self.batch_end = 0 + self.label_prop = label_prop + self.label_type = None + + find_image_count = 0 + for i in range(len(query)): + + name = list(query[i].keys())[0] + if name == "FindImage": + self.find_image_idx = i + find_image_count += 1 + + if find_image_count != 1: + logger.error( + "Query error. The query must contain exactly one FindImage command") + raise Exception('Query Error') + + if not "results" in self.query[self.find_image_idx]["FindImage"]: + self.query[self.find_image_idx]["FindImage"]["results"] = {} + + if self.label_prop is not None: + results = self.query[self.find_image_idx]["FindImage"]["results"] + if "list" not in results: + results["list"] = [] + if self.label_prop not in results["list"]: + results["list"].append(self.label_prop) + + self.query[self.find_image_idx]["FindImage"]["batch"] = {} + self.query[self.find_image_idx]["FindImage"]["blobs"] = False + + try: + _, r, b = execute_query( + client=self.client, query=self.query, blobs=[]) + batch = r[self.find_image_idx]["FindImage"]["batch"] + self.total_elements = batch["total_elements"] + except: + logger.error( + f"Query error: {self.query} {self.client.get_last_response_str()}") + raise + + self.query[self.find_image_idx]["FindImage"]["blobs"] = True + + def is_in_range(self, index): + + if index >= self.batch_start and index < self.batch_end: + return True + + return False + + def get_batch(self, index): + + total_batches = math.ceil(self.total_elements / self.batch_size) + batch_idx = math.floor(index / self.batch_size) + + if batch_idx >= total_batches: + raise Exception("Index out of range") + + query = self.query + qbatch = query[self.find_image_idx]["FindImage"]["batch"] + qbatch["batch_size"] = self.batch_size + qbatch["batch_id"] = batch_idx + + query[self.find_image_idx]["FindImage"]["batch"] = qbatch + + try: + + # This is to handle potential issues with + # disconnection/timeout and SSL context on multiprocessing + connection_ok = False + try: + _, r, b = execute_query( + query=self.query, blobs=[], client=self.client) + connection_ok = True + except: + # Connection failed, we retry just once to re-connect + self.client = self.client.clone() + + if not connection_ok: + # Connection failed, we have reconnected, we try again. + _, r, b = execute_query( + query=self.query, blobs=[], client=self.client) + + if len(b) == 0: + logger.error(f"index: {index}") + raise Exception("No results returned from ApertureDB") + + self.batch_images = b + self.batch_start = self.batch_size * batch_idx + self.batch_end = self.batch_start + len(b) + + if self.label_prop: + entities = r[self.find_image_idx]["FindImage"]["entities"] + try: + self.batch_labels = [l[self.label_prop] for l in entities] + except KeyError: + prop = self.label_prop + msg = ( + f"Property '{prop}' not found in some entities. " + "Ensure all entities have this property." + ) + logger.error(msg) + raise + else: + self.batch_labels = ["none" for l in range(len(b))] + except: + logger.error(f"Query error: {self.client.get_last_response_str()}") + raise + + def generator(self): + for index in range(self.total_elements): + if not self.is_in_range(index): + self.get_batch(index) + + idx = index % self.batch_size + img = self.batch_images[idx] + label = self.batch_labels[idx] + + nparr = np.frombuffer(img, dtype=np.uint8) + img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + yield img, label + + def get_dataset(self): + import tensorflow as tf + + if self.label_type is None: + if self.total_elements > 0: + self.get_batch(0) + if isinstance(self.batch_labels[0], int): + self.label_type = tf.int32 + elif isinstance(self.batch_labels[0], float): + self.label_type = tf.float32 + else: + self.label_type = tf.string + else: + self.label_type = tf.string + + return tf.data.Dataset.from_generator( + self.generator, + output_signature=( + tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), + tf.TensorSpec(shape=(), dtype=self.label_type) + ) + ) diff --git a/test/test_tf_connector.py b/test/test_tf_connector.py new file mode 100644 index 00000000..f7ac96c2 --- /dev/null +++ b/test/test_tf_connector.py @@ -0,0 +1,139 @@ +import time +import logging + +import tensorflow as tf +from aperturedb.TensorFlowDataset import ApertureDBTensorFlowDataset + +from aperturedb.ConnectorRest import ConnectorRest + +logger = logging.getLogger(__name__) + + +class TestTfDatasets(): + def validate_dataset(self, dataset: tf.data.Dataset, expected_length): + start = time.time() + + count = 0 + # Iterate over dataset. + for img, label in dataset: + if tf.size(img).numpy() == 0: + logger.error("Empty image?") + assert False + count += 1 + assert count == expected_length + + time_taken = time.time() - start + if time_taken != 0: + logger.info(f"Throughput (imgs/s): {expected_length / time_taken}") + + def test_nativeConstraints(self, db, utils, images): + assert len(images) > 0 + # This is a hack against a bug in batch API. + dim = 224 if isinstance(db, ConnectorRest) else 225 + query = [{ + "FindImage": { + "constraints": { + "age": [">=", 0] + }, + "operations": [ + { + "type": "resize", + "width": dim, + "height": dim + } + ], + "results": { + "list": ["license"] + } + } + }] + + dataset_wrapper = ApertureDBTensorFlowDataset( + db, query, label_prop="license") + dataset = dataset_wrapper.get_dataset() + + self.validate_dataset(dataset, utils.count_images()) + + def test_batchedDataset(self, db, utils, images): + len_limit = utils.count_images() + # This is a hack against a bug in batch API. + dim = 224 if isinstance(db, ConnectorRest) else 225 + query = [{ + "FindImage": { + "constraints": { + "age": [">=", 0] + }, + "operations": [ + { + "type": "resize", + "width": dim, + "height": dim + } + ], + "results": { + "list": ["license"], + "limit": len_limit + } + } + }] + + dataset_wrapper = ApertureDBTensorFlowDataset( + db, query, label_prop="license") + dataset = dataset_wrapper.get_dataset() + # Test batched dataset + batched_dataset = dataset.batch(10).prefetch(tf.data.AUTOTUNE) + + start = time.time() + count = 0 + for imgs, labels in batched_dataset: + count += tf.shape(imgs)[0].numpy() + assert count == len_limit + + time_taken = time.time() - start + if time_taken != 0: + logger.info(f"Throughput (imgs/s): {len_limit / time_taken}") + + def test_dynamic_label_dtype(self): + from unittest.mock import patch + import numpy as np + import cv2 + + class DummyClient: + def clone(self): + return self + + def get_last_response_str(self): + return "" + + query = [{"FindImage": {"results": {"list": ["prop"]}}}] + + with patch('aperturedb.TensorFlowDataset.execute_query') as mock_exec: + def side_effect_int(*args, **kwargs): + batch_dict = {"total_elements": 1} + entities = [{"prop": 42}] # int + r = [{"FindImage": {"batch": batch_dict, "entities": entities}}] + img = np.zeros((10, 10, 3), dtype=np.uint8) + _, b_img = cv2.imencode('.jpg', img) + b = [b_img.tobytes()] + return None, r, b + + mock_exec.side_effect = side_effect_int + dataset_wrapper = ApertureDBTensorFlowDataset( + DummyClient(), query, label_prop="prop") + dataset = dataset_wrapper.get_dataset() + assert dataset.element_spec[1].dtype == tf.int32 + + def side_effect_float(*args, **kwargs): + batch_dict = {"total_elements": 1} + entities = [{"prop": 3.14}] # float + r = [{"FindImage": {"batch": batch_dict, "entities": entities}}] + img = np.zeros((10, 10, 3), dtype=np.uint8) + _, b_img = cv2.imencode('.jpg', img) + b = [b_img.tobytes()] + return None, r, b + + mock_exec.side_effect = side_effect_float + dataset_wrapper_f = ApertureDBTensorFlowDataset(DummyClient( + ), [{"FindImage": {"results": {"list": ["prop"]}}}], label_prop="prop") + dataset_f = dataset_wrapper_f.get_dataset() + assert dataset_f.element_spec[1].dtype == tf.float32 diff --git a/test/test_torch_connector.py b/test/test_torch_connector.py index 2b3c5740..22dd91af 100644 --- a/test/test_torch_connector.py +++ b/test/test_torch_connector.py @@ -31,7 +31,7 @@ def validate_dataset(self, dataset: Union[DataLoader, Dataset], expected_length) if time_taken != 0: logger.info(f"Throughput (imgs/s): {len(dataset) / time_taken}") - def test_nativeContraints(self, db, utils, images): + def test_nativeConstraints(self, db, utils, images): assert len(images) > 0 # This is a hack against a bug in batch API. dim = 224 if isinstance(db, ConnectorRest) else 225 From 48ca416b774466f89ddfa06d1548f26a55841c47 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Fri, 22 May 2026 11:08:38 -0700 Subject: [PATCH 03/19] fix: optimize base64 decoding and reduce string allocations in ConnectorRest (#712) --- aperturedb/ConnectorRest.py | 11 +++++++++-- pyproject.toml | 3 +++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/aperturedb/ConnectorRest.py b/aperturedb/ConnectorRest.py index 803ab733..c3231234 100644 --- a/aperturedb/ConnectorRest.py +++ b/aperturedb/ConnectorRest.py @@ -33,6 +33,11 @@ from types import SimpleNamespace from typing import Optional +try: + import pybase64 as base64 +except ImportError: + import base64 + from aperturedb.Connector import Connector from aperturedb.Configuration import Configuration from requests.adapters import HTTPAdapter @@ -164,8 +169,10 @@ def _query(self, query, blob_array = [], try_resume=True): verify = self.config.use_ssl and self.config.verify_hostname) if response.status_code == 200: # Parse response: - json_response = json.loads(response.text) - import base64 + if hasattr(response, "json"): + json_response = response.json() + else: + json_response = json.loads(response.text) response_blob_array = [base64.b64decode( b) for b in json_response['blobs']] self.last_response = json_response["json"] diff --git a/pyproject.toml b/pyproject.toml index de914ca0..e6c8f2eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,9 @@ aperturedb = "aperturedb" "Bug Reports" = "https://github.com/aperture-data/aperturedb-python/issues" [project.optional-dependencies] +speedups = [ + "pybase64", +] # This is used when we build the docker image for notebook notebook = [ "torch", From 3391b0652f7172dceb5a1b420a4f06c3a632bdce Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Fri, 22 May 2026 18:01:10 -0700 Subject: [PATCH 04/19] docs: add examples tying the various data loaders and explain their hierarchy (#718) --- examples/README.md | 1 + examples/loaders_101/data_loader_hierarchy.py | 109 ++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 examples/loaders_101/data_loader_hierarchy.py diff --git a/examples/README.md b/examples/README.md index 5e4596b5..203af22d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -16,6 +16,7 @@ The following files are under *loaders_101* | File | Description | instructions | | -----| ------------| -----| | loaders.ipynb | A notebook with some sample code for aperturedb | Also available to read at [Aperturedb documentation](https://docs.aperturedata.io/HowToGuides/Advanced/loaders)| +| data_loader_hierarchy.py | A script explaining the data loader hierarchy and tying them together | ``python data_loader_hierarchy.py``| ## Example 2: Image classification using a pretrained model The following files are under *image_classification* diff --git a/examples/loaders_101/data_loader_hierarchy.py b/examples/loaders_101/data_loader_hierarchy.py new file mode 100644 index 00000000..ef0635f0 --- /dev/null +++ b/examples/loaders_101/data_loader_hierarchy.py @@ -0,0 +1,109 @@ +""" +Data Loader Hierarchy and Examples + +ApertureDB python SDK uses `ParallelLoader` as the main mechanism to efficiently batch and load data in parallel. +The `ParallelLoader` relies on the `Subscriptable` interface. +Classes that inherit from `Subscriptable` can be passed to `ParallelLoader.ingest()`. + +The main data loaders provided by ApertureDB handle parsing CSVs and generating the appropriate ApertureDB queries. +The hierarchy is as follows: + +Subscriptable + └── CSVParser + ├── EntityDataCSV (Loads entities and their properties) + ├── ImageDataCSV (Loads images as entities and uploads image blobs) + ├── VideoDataCSV (Loads videos) + ├── BlobDataCSV (Loads generic blobs) + ├── DescriptorDataCSV (Loads descriptors, e.g., embeddings) + ├── ConnectionDataCSV (Loads connections between entities) + ├── BBoxDataCSV (Loads bounding boxes) + └── PolygonDataCSV (Loads polygons) + +Other non-CSV data loaders (like PyTorchData, KaggleData) also inherit from `Subscriptable`. + +The following script demonstrates how to instantiate different loaders and tie them together using `ParallelLoader`. +""" + +import os +import base64 +import tempfile +import pandas as pd +from aperturedb.Connector import Connector +from aperturedb.ParallelLoader import ParallelLoader +from aperturedb.EntityDataCSV import EntityDataCSV +from aperturedb.ImageDataCSV import ImageDataCSV +from aperturedb.ConnectionDataCSV import ConnectionDataCSV + + +def create_sample_csvs(base_dir): + # 1. Create a sample CSV for Entities (Persons) + df_persons = pd.DataFrame({ + "EntityClass": ["Person", "Person"], + "name": ["Alice", "Bob"], + "age": [25, 30], + "constraint_name": ["Alice", "Bob"] + }) + persons_csv = os.path.join(base_dir, "persons.csv") + df_persons.to_csv(persons_csv, index=False) + + # 2. Create a sample CSV for Images + # Note: the paths should point to real images in a real scenario + dummy_image_path = os.path.join(base_dir, "dummy_image.png") + # 1x1 transparent PNG base64 + tiny_png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" + with open(dummy_image_path, "wb") as f: + f.write(base64.b64decode(tiny_png_b64)) + + df_images = pd.DataFrame({ + "filename": [dummy_image_path, dummy_image_path], + "image_id": ["img1", "img2"], + "source": ["camera1", "camera2"], + "constraint_image_id": ["img1", "img2"] + }) + images_csv = os.path.join(base_dir, "images.csv") + df_images.to_csv(images_csv, index=False) + + # 3. Create a sample CSV for Connections (Person -> Image) + df_connections = pd.DataFrame({ + "ConnectionClass": ["HasImage", "HasImage"], + "Person@name": ["Alice", "Bob"], + "_Image@image_id": ["img1", "img2"], + "connection_property": ["owns", "likes"] + }) + connections_csv = os.path.join(base_dir, "connections.csv") + df_connections.to_csv(connections_csv, index=False) + + return persons_csv, images_csv, connections_csv + + +def main(): + # Connect to ApertureDB (Make sure your ApertureDB instance is running) + db = Connector() + + # Initialize the ParallelLoader + # The ParallelLoader handles the actual ingestion of queries produced by the Data Loader objects + loader = ParallelLoader(db) + + with tempfile.TemporaryDirectory() as temp_dir: + persons_csv, images_csv, connections_csv = create_sample_csvs(temp_dir) + + # 1. Load Entities + print("Loading Entities...") + # Properties are derived directly from the CSV header names. + person_loader = EntityDataCSV(persons_csv) + loader.ingest(person_loader) + + # 2. Load Images + print("Loading Images...") + image_loader = ImageDataCSV(images_csv) + loader.ingest(image_loader) + + # 3. Load Connections + print("Loading Connections...") + connection_loader = ConnectionDataCSV(connections_csv) + loader.ingest(connection_loader) + print("Done loading data!") + + +if __name__ == "__main__": + main() From d5e13e6784280eee25a22266544c9cb666a06250 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Fri, 22 May 2026 20:00:39 -0700 Subject: [PATCH 05/19] fix: Update operations on DW Operations (#688) --- aperturedb/Images.py | 20 +++++++---- aperturedb/Operations.py | 48 +++++++++++++++++++++++--- test/test_Images.py | 25 ++++++++++++++ test/test_Operations.py | 74 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 157 insertions(+), 10 deletions(-) create mode 100644 test/test_Operations.py diff --git a/aperturedb/Images.py b/aperturedb/Images.py index 55d91d83..1c5cc3a0 100644 --- a/aperturedb/Images.py +++ b/aperturedb/Images.py @@ -69,7 +69,7 @@ def rotate(points, angle, c_x=0, c_y=0): points: The rotated points as a NumPy array of shape (n,2) and type int """ ANGLE = np.deg2rad(angle) - return np.array( + return np.round(np.array( [ [ c_x + np.cos(ANGLE) * (px - c_x) - np.sin(ANGLE) * (py - c_y), @@ -77,7 +77,7 @@ def rotate(points, angle, c_x=0, c_y=0): ] for px, py in points ] - ).astype(int) + )).astype(int) def resolve(points: np.array, image_meta, operations) -> np.array: @@ -100,12 +100,20 @@ def resolve(points: np.array, image_meta, operations) -> np.array: image_meta_height = image_meta["adb_image_height"] for operation in operations: if operation["type"] == "resize": - x_ratio = operation["width"] / image_meta_width - y_ratio = operation["height"] / image_meta_height + if "scale" in operation: + x_ratio = operation["scale"] + y_ratio = operation["scale"] + else: + x_ratio = operation["width"] / image_meta_width + y_ratio = operation["height"] / image_meta_height np.multiply(resolved, [x_ratio, y_ratio], out=resolved, casting="unsafe") - image_meta_width = operation["width"] - image_meta_height = operation["height"] + if "scale" in operation: + image_meta_width *= x_ratio + image_meta_height *= y_ratio + else: + image_meta_width = operation["width"] + image_meta_height = operation["height"] if operation["type"] == "rotate": angle = operation["angle"] resolved = rotate( diff --git a/aperturedb/Operations.py b/aperturedb/Operations.py index 88482dda..e96ca592 100644 --- a/aperturedb/Operations.py +++ b/aperturedb/Operations.py @@ -1,4 +1,5 @@ from __future__ import annotations +from typing import Optional class Operations(object): @@ -15,13 +16,25 @@ def __init__(self): def get_operations_arr(self): return self.operations_arr - def resize(self, width: int, height: int) -> Operations: + def resize(self, width: Optional[int] = None, height: Optional[int] = None, scale: Optional[float] = None) -> Operations: + + if scale is not None: + if width is not None or height is not None: + raise ValueError( + "Provide either 'scale' or both 'width' and 'height', but not a mix.") + else: + if width is None or height is None: + raise ValueError( + "Provide either 'scale' or both 'width' and 'height'.") op = { - "type": "resize", - "width": width, - "height": height, + "type": "resize" } + if width is not None: + op["width"] = width + op["height"] = height + if scale is not None: + op["scale"] = scale self.operations_arr.append(op) return self @@ -71,3 +84,30 @@ def interval(self, start: int, stop: int, step: int) -> Operations: self.operations_arr.append(op) return self + + def threshold(self, value: int) -> Operations: + + op = { + "type": "threshold", + "value": value, + } + + self.operations_arr.append(op) + return self + + def preview(self, max_frame_count: Optional[int] = None, max_time_fraction: Optional[float] = None, max_time_offset: Optional[str] = None, max_size_mb: Optional[float] = None) -> Operations: + + op = { + "type": "preview" + } + if max_frame_count is not None: + op["max_frame_count"] = max_frame_count + if max_time_fraction is not None: + op["max_time_fraction"] = max_time_fraction + if max_time_offset is not None: + op["max_time_offset"] = max_time_offset + if max_size_mb is not None: + op["max_size_mb"] = max_size_mb + + self.operations_arr.append(op) + return self diff --git a/test/test_Images.py b/test/test_Images.py index 20026e5b..7c41a26c 100644 --- a/test/test_Images.py +++ b/test/test_Images.py @@ -22,6 +22,17 @@ def test_resolve_resize(): assert resolved[1][1] == 10 +def test_resolve_resize_scale(): + points = np.array([[10, 10], [20, 20]], dtype=float) + meta = {"adb_image_width": 100, "adb_image_height": 100} + operations = [{"type": "resize", "scale": 0.5}] + resolved = resolve(points, meta, operations) + assert resolved[0][0] == 5 + assert resolved[0][1] == 5 + assert resolved[1][0] == 10 + assert resolved[1][1] == 10 + + def test_resolve_rotate(): points = np.array([[10, 10]], dtype=float) meta = {"adb_image_width": 100, "adb_image_height": 100} @@ -32,6 +43,20 @@ def test_resolve_rotate(): assert resolved[0][0] == 90 and abs(resolved[0][1] - 10) <= 1 +def test_resolve_ignored_operations(): + points = np.array([[10, 10], [20, 20]], dtype=float) + meta = {"adb_image_width": 100, "adb_image_height": 100} + operations = [ + {"type": "flip", "code": "horizontal"}, + {"type": "crop", "x": 10, "y": 10, "width": 50, "height": 50}, + {"type": "interval", "start": 0, "stop": 10, "step": 1}, + {"type": "threshold", "value": 128}, + {"type": "preview"} + ] + resolved = resolve(points, meta, operations) + assert np.array_equal(resolved, points) + + class MockClient: def __init__(self): self.responses = [] diff --git a/test/test_Operations.py b/test/test_Operations.py new file mode 100644 index 00000000..66a4907c --- /dev/null +++ b/test/test_Operations.py @@ -0,0 +1,74 @@ +import pytest +from aperturedb.Operations import Operations + + +class TestOperations: + def test_empty_operations(self): + op = Operations() + assert op.get_operations_arr() == [] + + def test_resize_width_height(self): + op = Operations().resize(width=100, height=200) + assert op.get_operations_arr() == [ + {"type": "resize", "width": 100, "height": 200}] + + def test_resize_scale(self): + op = Operations().resize(scale=0.5) + assert op.get_operations_arr() == [{"type": "resize", "scale": 0.5}] + + def test_rotate(self): + op = Operations().rotate(angle=90) + assert op.get_operations_arr() == [ + {"type": "rotate", "angle": 90, "resize": False}] + + def test_rotate_resize(self): + op = Operations().rotate(angle=90, resize=True) + assert op.get_operations_arr() == [ + {"type": "rotate", "angle": 90, "resize": True}] + + def test_flip(self): + op = Operations().flip(code="horizontal") + assert op.get_operations_arr() == [ + {"type": "flip", "code": "horizontal"}] + + def test_crop(self): + op = Operations().crop(x=10, y=20, width=100, height=200) + assert op.get_operations_arr() == [ + {"type": "crop", "x": 10, "y": 20, "width": 100, "height": 200}] + + def test_interval(self): + op = Operations().interval(start=0, stop=100, step=2) + assert op.get_operations_arr() == [ + {"type": "interval", "start": 0, "stop": 100, "step": 2}] + + def test_threshold(self): + op = Operations().threshold(value=128) + assert op.get_operations_arr() == [{"type": "threshold", "value": 128}] + + def test_preview(self): + op = Operations().preview(max_frame_count=10, max_time_fraction=0.5) + assert op.get_operations_arr() == [ + {"type": "preview", "max_frame_count": 10, "max_time_fraction": 0.5}] + + def test_preview_all_args(self): + op = Operations().preview(max_frame_count=10, max_time_fraction=0.5, + max_time_offset="00:00:10", max_size_mb=10.5) + assert op.get_operations_arr() == [ + {"type": "preview", "max_frame_count": 10, "max_time_fraction": 0.5, "max_time_offset": "00:00:10", "max_size_mb": 10.5}] + + def test_chained_operations(self): + op = Operations().resize(width=100, height=100).rotate( + angle=90).crop(x=0, y=0, width=50, height=50) + assert op.get_operations_arr() == [ + {"type": "resize", "width": 100, "height": 100}, + {"type": "rotate", "angle": 90, "resize": False}, + {"type": "crop", "x": 0, "y": 0, "width": 50, "height": 50} + ] + + def test_resize_invalid_args(self): + with pytest.raises(ValueError, match="Provide either 'scale' or both 'width' and 'height'"): + Operations().resize() + with pytest.raises(ValueError, match="Provide either 'scale' or both 'width' and 'height'"): + Operations().resize(width=100) + with pytest.raises(ValueError, match="Provide either 'scale' or both 'width' and 'height', but not a mix"): + Operations().resize(width=100, height=100, scale=0.5) From a3e393134f9265b1b9fc28f47bfb918d1d5efa38 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sat, 23 May 2026 09:05:43 -0700 Subject: [PATCH 06/19] fix: check all statuses in check_status to report errors properly (#730) --- aperturedb/Connector.py | 31 +++++++++++++++++++------------ test/test_Connector.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 12 deletions(-) create mode 100644 test/test_Connector.py diff --git a/aperturedb/Connector.py b/aperturedb/Connector.py index 7baa0395..1bd71b42 100644 --- a/aperturedb/Connector.py +++ b/aperturedb/Connector.py @@ -25,7 +25,7 @@ # THE SOFTWARE. # from __future__ import annotations -from typing import Optional +from typing import Optional, Any from . import queryMessage import sys import os @@ -43,7 +43,6 @@ from types import SimpleNamespace from dataclasses import dataclass from aperturedb.Configuration import Configuration -from aperturedb.types import CommandResponses from aperturedb.LoggingUtils import censor_tokens logger = logging.getLogger(__name__) @@ -669,28 +668,36 @@ def last_query_ok(self): return self.check_status(self.response) >= 0 - def check_status(self, json_res: CommandResponses) -> int: + def check_status(self, json_res: Any) -> int: """ - Returns the status of the first command response from the server. - Can traverse a JSON recursively to find the first status. + Returns the status of the first negative command response from the server, + or the status of the first command if all are non-negative. + Can traverse a JSON recursively to find the statuses. Args: - json_res (CommandResponses): The actual response from the server. + json_res (Any): The actual response from the server. Returns: - int: The value recieved from the server, or -2 if not found. + int: The value received from the server, or -2 if not found. """ # Default status is -2, which is an error, but not a server error. status = STATUS_ERROR_DEFAULT if (isinstance(json_res, dict)): if ("status" not in json_res): - status = self.check_status(json_res[list(json_res.keys())[0]]) + for i, val in enumerate(json_res.values()): + st = self.check_status(val) + if i == 0: + status = st + if st < 0: + return st else: status = json_res["status"] elif (isinstance(json_res, (tuple, list))): - if ("status" not in json_res[0]): - status = self.check_status(json_res[0]) - else: - status = json_res[0]["status"] + for i, res in enumerate(json_res): + st = self.check_status(res) + if i == 0: + status = st + if st < 0: + return st return status diff --git a/test/test_Connector.py b/test/test_Connector.py new file mode 100644 index 00000000..0d5cd45b --- /dev/null +++ b/test/test_Connector.py @@ -0,0 +1,31 @@ +from aperturedb.Connector import Connector + + +class TestConnector: + def test_check_status_mixed(self): + # We don't need a real connection, just the check_status method + connector = Connector(connect=False) + + # Test case 1: all ok + res1 = [{"FindImage": {"status": 0}}, {"FindImage": {"status": 0}}] + assert connector.check_status(res1) == 0 + + # Test case 2: first is ok, second is negative + res2 = [{"FindImage": {"status": 0}}, {"FindImage": {"status": -4}}] + assert connector.check_status(res2) == -4 + + # Test case 3: first is negative, second is ok + res3 = [{"FindImage": {"status": -1}}, {"FindImage": {"status": 0}}] + assert connector.check_status(res3) == -1 + + # Test case 4: deeply nested dict + res4 = {"FindImage": {"status": 0}} + assert connector.check_status(res4) == 0 + + # Test case 5: nested ok, nested negative + res5 = [{"AddImage": {"status": 0}}, {"AddEntity": {"status": -3}}] + assert connector.check_status(res5) == -3 + + # Test case 6: multiple keys in a single dict (tests values traversal) + res6 = {"FindImage": {"status": 0}, "FindEntity": {"status": -4}} + assert connector.check_status(res6) == -4 From 3bbc9a8d5a53467f3fd40a02b88cbd0e0c257daa Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sat, 23 May 2026 10:03:34 -0700 Subject: [PATCH 07/19] Move generateImage to adb to allow generalized usage (#719) --- .../cli/generate_images.py | 0 aperturedb/cli/utilities.py | 46 ++++++++++++++++++- test/generateInput.py | 2 +- test/run_test_container.sh | 26 +++++++++-- 4 files changed, 69 insertions(+), 5 deletions(-) rename test/generateImages.py => aperturedb/cli/generate_images.py (100%) diff --git a/test/generateImages.py b/aperturedb/cli/generate_images.py similarity index 100% rename from test/generateImages.py rename to aperturedb/cli/generate_images.py diff --git a/aperturedb/cli/utilities.py b/aperturedb/cli/utilities.py index c5b6a189..51baa884 100644 --- a/aperturedb/cli/utilities.py +++ b/aperturedb/cli/utilities.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Annotated +from typing import Annotated, Optional import typer @@ -97,3 +97,47 @@ def visualize_schema( s = utils.visualize_schema() result = s.render(filename, format=format) print(result) + + +class ImageType(str, Enum): + PNG = "png" + JPG = "jpg" + + +@app.command(name="generate-images", help="Generate placeholder images") +def generate_images( + count: int = typer.Option( + 1, "--count", "-c", help="Number of images to generate"), + size: str = typer.Option("256x256", "--size", "-s", + help="Size of images (widthxheight)"), + output: str = typer.Option( + "/tmp/generated_image%%", "--output", "-o", help="Output file path pattern"), + zerofill: int = typer.Option( + 0, "--zerofill", "-z", help="Number of zeros to pad the index with"), + imagetype: ImageType = typer.Option( + ImageType.PNG, "--imagetype", "-t", help="Type of image to generate (png, jpg)"), + manifest: Optional[str] = typer.Option( + None, "--manifest", "-m", help="Manifest file to write generated image paths to"), + append_text: str = typer.Option( + "", "--append-text", "-a", help="Text to append to the image") +): + import aperturedb.cli.generate_images as generate_images_mod + + size_tuple = (256, 256) + try: + parsed_size = generate_images_mod.ImageSize.parse(size) + size_tuple = (parsed_size.width, parsed_size.height) + except Exception as e: + typer.echo(f"Invalid size: {e}") + raise typer.Abort() + + generator = generate_images_mod.ImageGenerator( + count=count, + size=size_tuple, + output=output, + zerofill=zerofill, + image_type=imagetype.value, + manifest=manifest, + append_text=append_text + ) + generator.run() diff --git a/test/generateInput.py b/test/generateInput.py index b1bc6e06..d1f90f0b 100644 --- a/test/generateInput.py +++ b/test/generateInput.py @@ -7,7 +7,7 @@ import pandas as pd import numpy as np -from generateImages import ImageGenerator +from aperturedb.cli.generate_images import ImageGenerator from itertools import product diff --git a/test/run_test_container.sh b/test/run_test_container.sh index fd537137..4904c348 100755 --- a/test/run_test_container.sh +++ b/test/run_test_container.sh @@ -71,11 +71,31 @@ wait_for_stack() { local elapsed=0 echo "Waiting for stack ${tag} to become ready (timeout ${timeout}s)..." while [ $elapsed -lt $timeout ]; do + local lenz_ready=0 + local nginx_ready=0 + if docker run --rm --network=${network} curlimages/curl:latest \ - -sS -o /dev/null -m 2 http://lenz:58085/ >/dev/null 2>&1; then - echo "Stack ${tag} is ready after ${elapsed}s" - return 0 + nc -z -w 2 lenz 58085 >/dev/null 2>&1; then + lenz_ready=1 + fi + + if [[ "$tag" == *"_http" ]]; then + if docker run --rm --network=${network} curlimages/curl:latest \ + -sS -o /dev/null -m 2 http://nginx:80/ >/dev/null 2>&1; then + nginx_ready=1 + fi + + if [ "$lenz_ready" -eq 1 ] && [ "$nginx_ready" -eq 1 ]; then + echo "Stack ${tag} is ready after ${elapsed}s" + return 0 + fi + else + if [ "$lenz_ready" -eq 1 ]; then + echo "Stack ${tag} is ready after ${elapsed}s" + return 0 + fi fi + sleep 2 elapsed=$((elapsed + 2)) done From 2800b44f06ef48d78f3d726dea236e71e44a8833 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sat, 23 May 2026 17:46:45 -0700 Subject: [PATCH 08/19] fix: resolve conda_gpu tests (#722) --- .github/workflows/pr.yaml | 48 ++++++++++++++++++++++++++++++++++- docker/pytorch-gpu/Dockerfile | 3 ++- docker/pytorch-gpu/build.sh | 8 +++++- test/run_test_container.sh | 39 +++++++++++++++++++++++----- 4 files changed, 89 insertions(+), 9 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 6b4b3ddb..cc1af664 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -17,7 +17,7 @@ jobs: steps: - name: Cleanup previous run - run: docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb/db* /workspace/test/aperturedb/logs*" + run: docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb" continue-on-error: true - uses: actions/checkout@v3 @@ -67,3 +67,49 @@ jobs: LENZ_TAG: dev run: ./ci.sh shell: bash + + run_test_conda_gpu: + + runs-on: + - self-hosted + - deployer + + steps: + + - name: Cleanup previous run + run: docker run --rm -v ${{ github.workspace }}:/workspace alpine sh -c "rm -rf /workspace/test/aperturedb" + continue-on-error: true + + - uses: actions/checkout@v3 + + - name: Login to DockerHub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKER_USER }} + password: ${{ secrets.DOCKER_PASS }} + + - name: Login to Google Cloud + uses: google-github-actions/setup-gcloud@v0 + with: + service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + export_default_credentials: true + + - name: Build tests on pytorch GPU image + run: | + rm -rf docker/pytorch-gpu/aperturedata + mkdir -p docker/pytorch-gpu/aperturedata + cp -r aperturedb pyproject.toml README.md test docker/pytorch-gpu/aperturedata + bash docker/pytorch-gpu/build.sh + shell: bash + + - name: Run Tests + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_DEFAULT_REGION: ${{ secrets.AWS_DEFAULT_REGION }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + GCP_SERVICE_ACCOUNT_KEY: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + RUNNER_NAME: ${{ runner.name }}_${{ github.run_id }} + run: | + cd test + ./run_test_container.sh aperturedata/aperturedb-pytorch-gpu diff --git a/docker/pytorch-gpu/Dockerfile b/docker/pytorch-gpu/Dockerfile index 846c7fc6..f2ac062c 100644 --- a/docker/pytorch-gpu/Dockerfile +++ b/docker/pytorch-gpu/Dockerfile @@ -1,4 +1,4 @@ -FROM gcr.io/deeplearning-platform-release/pytorch-gpu.1-13.py37 +FROM gcr.io/deeplearning-platform-release/pytorch-gpu.py310:m130 RUN mkdir /aperturedata ADD docker/pytorch-gpu/aperturedata /aperturedata @@ -6,6 +6,7 @@ ADD docker/pytorch-gpu/aperturedata /aperturedata RUN pip install awscli RUN apt-get update && apt-get install -y libopencv-dev python3-opencv fuse libfuse-dev RUN cd /aperturedata && pip install -e ".[dev]" +RUN pip install git+https://github.com/openai/CLIP.git@d05afc436d78f1c48dc0dbf8e5980a9d471f35f6 COPY docker/pytorch-gpu/scripts/start.sh /start.sh RUN chmod 755 /start.sh diff --git a/docker/pytorch-gpu/build.sh b/docker/pytorch-gpu/build.sh index 04d9c511..1ec27642 100644 --- a/docker/pytorch-gpu/build.sh +++ b/docker/pytorch-gpu/build.sh @@ -1 +1,7 @@ -docker build -f docker/pytorch-gpu/Dockerfile -t aperturedata/aperturedb-pytorch-gpu:latest . \ No newline at end of file +if [ "${NO_CACHE:-false}" = "true" ]; then + CACHE_FLAG="--no-cache" +else + CACHE_FLAG="" +fi + +docker build --pull $CACHE_FLAG -f docker/pytorch-gpu/Dockerfile -t aperturedata/aperturedb-pytorch-gpu:latest . \ No newline at end of file diff --git a/test/run_test_container.sh b/test/run_test_container.sh index 4904c348..aaa8900d 100755 --- a/test/run_test_container.sh +++ b/test/run_test_container.sh @@ -9,17 +9,26 @@ function check_containers_networks(){ docker network ls } +function get_sudo() { + if command -v sudo >/dev/null 2>&1 && sudo -n true 2>/dev/null; then + echo "sudo" + else + echo "" + fi +} + function run_aperturedb_instance(){ set -e TAG=$1 #Ensure clean environment (as much as possible) RUNNER_NAME=$TAG docker compose -f docker-compose.yml down --remove-orphans docker network rm ${TAG}_host_default || true + docker network rm ${TAG}_default || true # ensure latest db docker compose pull - rm -rf output + $(get_sudo) rm -rf output mkdir -m 777 output docker network create ${TAG}_host_default @@ -37,6 +46,29 @@ function run_aperturedb_instance(){ IP_REGEX='[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}' +function teardown() { + echo "Tearing down containers and networks..." + if [ "$TEST_PROTOCOL" == "http" ] || [ "$TEST_PROTOCOL" == "both" ]; then + RUNNER_NAME="${RUNNER_NAME}_http" docker compose -f docker-compose.yml down --remove-orphans || true + docker network rm "${RUNNER_NAME}_http_host_default" || true + docker network rm "${RUNNER_NAME}_http_default" || true + fi + if [ "$TEST_PROTOCOL" == "non_http" ] || [ "$TEST_PROTOCOL" == "both" ]; then + RUNNER_NAME="${RUNNER_NAME}_non_http" docker compose -f docker-compose.yml down --remove-orphans || true + docker network rm "${RUNNER_NAME}_non_http_host_default" || true + docker network rm "${RUNNER_NAME}_non_http_default" || true + fi +} +trap teardown EXIT + +# The LOG_PATH and RUNNER_INFO_PATH are set to the current working directory +LOG_PATH="$(pwd)/aperturedb/logs" +TESTING_LOG_PATH="/aperturedb/test/server_logs" +RUNNER_INFO_PATH="$(pwd)/aperturedb/logs/runner_state" + +$(get_sudo) mkdir -p "$RUNNER_INFO_PATH" +$(get_sudo) chmod -R 777 "$LOG_PATH" || true + # Check if TEST_PROTOCOL is set, otherwise default to both TEST_PROTOCOL=${TEST_PROTOCOL:-"both"} @@ -48,11 +80,6 @@ if [ "$TEST_PROTOCOL" == "non_http" ] || [ "$TEST_PROTOCOL" == "both" ]; then GATEWAY_NON_HTTP=$(run_aperturedb_instance "${RUNNER_NAME}_non_http" | grep $IP_REGEX ) fi -# The LOG_PATH and RUNNER_INFO_PATH are set to the current working directory -LOG_PATH="$(pwd)/aperturedb/logs" -TESTING_LOG_PATH="/aperturedb/test/server_logs" -RUNNER_INFO_PATH="$(pwd)/aperturedb/logs/runner_state" - check_containers_networks | tee "$RUNNER_INFO_PATH"/runner_state.log REPOSITORY="aperturedata/aperturedb-python-tests" From ef926b91b3acf8e5322c016814e6a657fbce1b1c Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sat, 23 May 2026 18:32:15 -0700 Subject: [PATCH 09/19] fix: resolve unhashable type slice in ParallelQuery (#714) --- aperturedb/CommonLibrary.py | 3 +-- aperturedb/ParallelLoader.py | 8 +++--- aperturedb/ParallelQuery.py | 49 ++++++++++++++++++++---------------- test/test_Parallel.py | 36 ++++++++++++++++++++++++-- 4 files changed, 67 insertions(+), 29 deletions(-) diff --git a/aperturedb/CommonLibrary.py b/aperturedb/CommonLibrary.py index 6022349d..e147a0a3 100644 --- a/aperturedb/CommonLibrary.py +++ b/aperturedb/CommonLibrary.py @@ -375,8 +375,7 @@ def map_response_to_handler(handler, query, query_blobs, response, response_blo handler( query[start:end], query_blobs[blobs_start:blobs_end], - response[start:end] if isinstance( - response, list) else response, + response[start:end] if is_list else response, response_blobs[blobs_returned:blobs_returned + b_count] if len(response_blobs) >= blobs_returned + b_count else None, None if cmd_index_offset is None else cmd_index_offset + i) diff --git a/aperturedb/ParallelLoader.py b/aperturedb/ParallelLoader.py index 94344411..0d7efb30 100644 --- a/aperturedb/ParallelLoader.py +++ b/aperturedb/ParallelLoader.py @@ -203,10 +203,10 @@ def print_stats(self) -> None: print(f"Total time (s): {self.total_actions_time}") print(f"Total queries executed: {total_queries_exec}") - succeeded_queries = sum([stat["succeeded_queries"] - for stat in self.actual_stats]) - succeeded_commands = sum([stat["succeeded_commands"] - for stat in self.actual_stats]) + succeeded_queries = sum(stat["succeeded_queries"] + for stat in self.actual_stats) + succeeded_commands = sum(stat["succeeded_commands"] + for stat in self.actual_stats) if succeeded_queries == 0: print("All queries failed!") diff --git a/aperturedb/ParallelQuery.py b/aperturedb/ParallelQuery.py index 2e567bda..498aed7e 100644 --- a/aperturedb/ParallelQuery.py +++ b/aperturedb/ParallelQuery.py @@ -173,6 +173,7 @@ def process_responses(requests, input_blobs, responses, output_blobs): f"expected 6 > args > 3, got {parameter_count}") if parameter_count == 4: indexless_handler = response_handler + def response_handler(query, qblobs, resp, rblobs, qindex): return indexless_handler( query, qblobs, resp, rblobs) @@ -191,28 +192,34 @@ def response_handler(query, qblobs, resp, rblobs, qindex): return indexless_hand worker_stats["succeeded_commands"] = len(q) worker_stats["succeeded_queries"] = len(data) worker_stats["objects_existed"] = sum( - [v['status'] == 2 for i in r for k, v in i.items()]) + v['status'] == 2 for i in r for k, v in i.items()) elif result == 1: self.error_counter += 1 worker_stats["succeeded_queries"] = 0 worker_stats["succeeded_commands"] = 0 worker_stats["objects_existed"] = 0 elif result == 2: + if client.last_query_ok(): + query_time = client.get_last_query_time() # with result 2, some queries might have failed. - def filter_per_group(group): - return group.items() if isinstance(group, dict) else {} - worker_stats["succeeded_commands"] = sum( - [v['status'] == 0 for i in r for k, v in filter_per_group(i)]) - worker_stats["objects_existed"] = sum( - [v['status'] == 2 for i in r for k, v in filter_per_group(i)]) - sq = 0 - for i in range(0, len(r), self.commands_per_query): - # Some errors stop the whole query from being executed - # https://docs.aperturedata.io/query_language/Overview/Responses#return-status - if issubclass(type(r), list): - if all([v['status'] == 0 for j in r[i:i + self.commands_per_query] for k, v in filter_per_group(j)]): + if isinstance(r, list): + def filter_per_group(group): + return group.items() if isinstance(group, dict) else () + worker_stats["succeeded_commands"] = sum( + v['status'] == 0 for i in r for k, v in filter_per_group(i)) + worker_stats["objects_existed"] = sum( + v['status'] == 2 for i in r for k, v in filter_per_group(i)) + sq = 0 + for i in range(0, len(r), self.commands_per_query): + # Some errors stop the whole query from being executed + # https://docs.aperturedata.io/query_language/Overview/Responses#return-status + if all(v['status'] == 0 for j in r[i:i + self.commands_per_query] for k, v in filter_per_group(j)): sq += 1 - worker_stats["succeeded_queries"] = sq + worker_stats["succeeded_queries"] = sq + else: + worker_stats["succeeded_commands"] = 0 + worker_stats["objects_existed"] = 0 + worker_stats["succeeded_queries"] = 0 else: query_time = 1 worker_stats["succeeded_commands"] = len(q) @@ -254,16 +261,16 @@ def worker(self, thid: int, generator, start: int, end: int, run_event) -> None: logger.info(f"Worker {thid} executed {total_batches} batches") def get_objects_existed(self) -> int: - return sum([stat["objects_existed"] - for stat in self.actual_stats]) + return sum(stat["objects_existed"] + for stat in self.actual_stats) def get_succeeded_queries(self) -> int: - return sum([stat["succeeded_queries"] - for stat in self.actual_stats]) + return sum(stat["succeeded_queries"] + for stat in self.actual_stats) def get_succeeded_commands(self) -> int: - return sum([stat["succeeded_commands"] - for stat in self.actual_stats]) + return sum(stat["succeeded_commands"] + for stat in self.actual_stats) def query(self, generator, batchsize: int = 1, numthreads: int = 4, stats: bool = False) -> None: """ @@ -302,7 +309,7 @@ def query(self, generator, batchsize: int = 1, numthreads: int = 4, stats: bool self.print_stats() else: # allow subclass to do verification - if issubclass(type(self), ParallelQuery) and hasattr(self, 'verify_generator') and callable(self.verify_generator): + if isinstance(self, ParallelQuery) and hasattr(self, 'verify_generator') and callable(self.verify_generator): self.verify_generator(generator) elif len(generator) > 0: if isinstance(generator[0], tuple) and isinstance(generator[0][0], list): diff --git a/test/test_Parallel.py b/test/test_Parallel.py index f83cff03..85c7cd51 100644 --- a/test/test_Parallel.py +++ b/test/test_Parallel.py @@ -62,7 +62,7 @@ def test_someBadQueries(self, db: Connector): except Exception as e: print(e) print("Failed to renew Session") - assert False + raise def test_allBadQueries(self, db: Connector): """ @@ -80,7 +80,39 @@ def test_allBadQueries(self, db: Connector): except Exception as e: print(e) print("Failed to renew Session") - assert False + raise + + def test_dictResponseHandling(self): + """ + Verifies that it handles a dict response from a failing server properly. + Guards against regression to "unhashable type: 'slice'" when r is a dict. + """ + from unittest.mock import MagicMock + try: + elements = 10 + generator = GeneratorWithErrors(elements=elements) + db = MagicMock(spec=Connector) + db.clone.return_value = db + db.config = "mock_config" + db.query.return_value = ([{"GetSchema": {"status": 0}}], []) + querier = ParallelQuery(db, dry_run=False) + + # Now set the mock for the actual queries + db.query.return_value = ( + {"status": 3, "error_msg": "mock error"}, []) + db.last_query_ok.return_value = True + db.get_last_query_time.return_value = 1.0 + + querier.query(generator, batchsize=2, numthreads=1, stats=True) + + # Since all queries got status 3, no commands or queries succeeded. + assert querier.get_succeeded_commands() == 0 + assert querier.get_succeeded_queries() == 0 + # Ensure stats were recorded properly + assert len(querier.actual_stats) > 0 + except Exception as e: + print(e) + raise def test_dask_dry_run(db: Connector): From d26470d87d0569b97b216e6941a4360429509f39 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sat, 23 May 2026 20:54:06 -0700 Subject: [PATCH 10/19] fix(cli): reduce traceback length by disabling Typer pretty exceptions (#704) --- aperturedb/cli/adb.py | 3 ++- aperturedb/cli/configure.py | 2 +- aperturedb/cli/ingest.py | 2 +- aperturedb/cli/keys.py | 2 +- aperturedb/cli/tokens.py | 2 +- aperturedb/cli/transact.py | 2 +- aperturedb/cli/utilities.py | 2 +- 7 files changed, 8 insertions(+), 7 deletions(-) diff --git a/aperturedb/cli/adb.py b/aperturedb/cli/adb.py index 5bac0cfd..eeb42222 100644 --- a/aperturedb/cli/adb.py +++ b/aperturedb/cli/adb.py @@ -2,7 +2,8 @@ from aperturedb.cli import configure, ingest, utilities, transact -app = typer.Typer(pretty_exceptions_show_locals=False) +app = typer.Typer(pretty_exceptions_show_locals=False, + pretty_exceptions_enable=False) app.add_typer(ingest.app, name="ingest", help="Ingest data into ApertureDB.") app.add_typer(configure.app, name="config", diff --git a/aperturedb/cli/configure.py b/aperturedb/cli/configure.py index 5afd2ba6..db471cd6 100644 --- a/aperturedb/cli/configure.py +++ b/aperturedb/cli/configure.py @@ -32,7 +32,7 @@ def default(self, o): APP_NAME = "aperturedb" APP_NAME_CLI = "adb" -app = typer.Typer() +app = typer.Typer(pretty_exceptions_enable=False) def _config_file_path(as_global: bool) -> Path: diff --git a/aperturedb/cli/ingest.py b/aperturedb/cli/ingest.py index 4e7e665c..fc953056 100644 --- a/aperturedb/cli/ingest.py +++ b/aperturedb/cli/ingest.py @@ -14,7 +14,7 @@ from tqdm import tqdm logger = logging.getLogger(__file__) -app = typer.Typer() +app = typer.Typer(pretty_exceptions_enable=False) IngestType = Enum('IngestType', {k: str(k) for k in ObjectType._member_names_}) diff --git a/aperturedb/cli/keys.py b/aperturedb/cli/keys.py index adaa8cd3..faff7ee6 100644 --- a/aperturedb/cli/keys.py +++ b/aperturedb/cli/keys.py @@ -6,7 +6,7 @@ from aperturedb.Configuration import Configuration from aperturedb.Connector import Connector -app = typer.Typer() +app = typer.Typer(pretty_exceptions_enable=False) @app.command(help="Create Key for a user") diff --git a/aperturedb/cli/tokens.py b/aperturedb/cli/tokens.py index 88486d30..6ca10e6b 100644 --- a/aperturedb/cli/tokens.py +++ b/aperturedb/cli/tokens.py @@ -7,7 +7,7 @@ from aperturedb.CommonLibrary import create_connector, execute_query from aperturedb.Utils import Utils -app = typer.Typer() +app = typer.Typer(pretty_exceptions_enable=False) @app.command(help="List User Authentication Tokens") diff --git a/aperturedb/cli/transact.py b/aperturedb/cli/transact.py index 2684d0d7..416d23f3 100644 --- a/aperturedb/cli/transact.py +++ b/aperturedb/cli/transact.py @@ -26,7 +26,7 @@ def load_fuse(): "fuse not found for this env. This is not critical for adb to continue.") -app = typer.Typer(callback=load_fuse) +app = typer.Typer(callback=load_fuse, pretty_exceptions_enable=False) class OutputTypes(str, Enum): diff --git a/aperturedb/cli/utilities.py b/aperturedb/cli/utilities.py index 51baa884..8c2bad44 100644 --- a/aperturedb/cli/utilities.py +++ b/aperturedb/cli/utilities.py @@ -5,7 +5,7 @@ from aperturedb.cli.console import console -app = typer.Typer() +app = typer.Typer(pretty_exceptions_enable=False) import aperturedb.cli.keys as keys import aperturedb.cli.tokens as tokens From 12cbe4d9aa3d83eebb1dac2db5c5e7cc979319d2 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sat, 23 May 2026 22:58:53 -0700 Subject: [PATCH 11/19] Closes #739: Skip slow and external network tests in PR CI (#740) --- .github/workflows/nightly.yml | 56 +++++++++++++++++++++++++++++++++++ .github/workflows/pr.yaml | 1 + test/run_test.sh | 10 ++++++- test/run_test_container.sh | 2 ++ 4 files changed, 68 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/nightly.yml diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml new file mode 100644 index 00000000..bb32ee55 --- /dev/null +++ b/.github/workflows/nightly.yml @@ -0,0 +1,56 @@ +name: nightly + +on: + schedule: + - cron: '0 0 * * *' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + build-test: + runs-on: + - benchmark + + steps: + + - uses: actions/checkout@v3 + + - name: Login to DockerHub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKER_USER }} + password: ${{ secrets.DOCKER_PASS }} + + - name: Login to Google Cloud + uses: google-github-actions/setup-gcloud@v0 + with: + service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + export_default_credentials: true + + - name: Build and Run Tests + env: + # Run both protocols inside a single job. run_test_container.sh starts + # the http and non_http stacks in parallel, so wall-clock time is + # roughly max(http, non_http) instead of the previous matrix approach + # which duplicated every docker image build (notebook deps / notebook + # / tests / coverage) on a second runner. + TEST_PROTOCOL: both + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_DEFAULT_REGION: ${{ secrets.AWS_DEFAULT_REGION }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + GCP_SERVICE_ACCOUNT_KEY: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + RUNNER_NAME: ${{ runner.name }}_nightly + # BuildKit + inline cache so subsequent runs actually reuse layers + # pulled via --cache-from in ci.sh. + DOCKER_BUILDKIT: 1 + COMPOSE_DOCKER_CLI_BUILD: 1 + ADB_REPO: aperturedata/aperturedb + ADB_TAG: dev + LENZ_REPO: aperturedata/lenz + LENZ_TAG: dev + run: RUN_TESTS=true ./ci.sh + shell: bash + diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index cc1af664..4f69a3a5 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -65,6 +65,7 @@ jobs: ADB_TAG: dev LENZ_REPO: aperturedata/lenz LENZ_TAG: dev + SKIP_SLOW_TESTS: true run: ./ci.sh shell: bash diff --git a/test/run_test.sh b/test/run_test.sh index 4bf1a21b..58d68a15 100755 --- a/test/run_test.sh +++ b/test/run_test.sh @@ -30,7 +30,15 @@ set +e CLIENT_PATH="${APERTUREDB_LOG_PATH}/../client/${FILTER}" CLIENT_PATH=${CLIENT_PATH// /_} mkdir -p ${CLIENT_PATH} -PROJECT=aperturedata KAGGLE_username=ci KAGGLE_key=dummy python3 -m pytest --cov=aperturedb -m "$FILTER" test_*.py -v | tee ${CLIENT_PATH}/test.log + +if [ "${SKIP_SLOW_TESTS:-false}" == "true" ]; then + echo "Skipping slow and external tests for this run..." + pytest_filter="${FILTER} and not slow and not external_network" +else + pytest_filter="${FILTER}" +fi + +PROJECT=aperturedata KAGGLE_username=ci KAGGLE_key=dummy python3 -m pytest --cov=aperturedb -m "$pytest_filter" test_*.py -v | tee ${CLIENT_PATH}/test.log RESULT=$? cp error*.log -v ${CLIENT_PATH} || true diff --git a/test/run_test_container.sh b/test/run_test_container.sh index aaa8900d..6ea9b0e0 100755 --- a/test/run_test_container.sh +++ b/test/run_test_container.sh @@ -155,6 +155,7 @@ if [ "$TEST_PROTOCOL" == "http" ] || [ "$TEST_PROTOCOL" == "both" ]; then -e APERTUREDB_LOG_PATH="${TESTING_LOG_PATH}" \ -e GATEWAY="nginx" \ -e FILTER="http" \ + -e SKIP_SLOW_TESTS="${SKIP_SLOW_TESTS:-false}" \ $REPOSITORY & pid1=$! fi @@ -174,6 +175,7 @@ if [ "$TEST_PROTOCOL" == "non_http" ] || [ "$TEST_PROTOCOL" == "both" ]; then -e APERTUREDB_LOG_PATH="${TESTING_LOG_PATH}" \ -e GATEWAY="lenz" \ -e FILTER="not http" \ + -e SKIP_SLOW_TESTS="${SKIP_SLOW_TESTS:-false}" \ $REPOSITORY & pid2=$! fi From 767cc8f33b897c6c6fe400bc48dffc0c40636076 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sun, 24 May 2026 10:45:05 -0700 Subject: [PATCH 12/19] Change ParallelQuerySet to use QueryBuilder (#725) --- .github/workflows/develop.yml | 12 +++++++----- .github/workflows/main.yml | 12 +++++++----- .github/workflows/pr.yaml | 12 +++++++----- .github/workflows/release.yaml | 12 +++++++----- aperturedb/ParallelQuerySet.py | 19 ++++++++---------- aperturedb/Query.py | 21 ++++++++++++++++++++ test/test_Datawizard.py | 35 ++++++++++++++++++++++++++++++++++ 7 files changed, 92 insertions(+), 31 deletions(-) diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml index db6d5394..295f4d26 100644 --- a/.github/workflows/develop.yml +++ b/.github/workflows/develop.yml @@ -28,12 +28,14 @@ jobs: username: ${{ secrets.DOCKER_USER }} password: ${{ secrets.DOCKER_PASS }} - - name: Login to Google Cloud - uses: google-github-actions/setup-gcloud@v0 + - name: Authenticate to Google Cloud + uses: google-github-actions/auth@v2 with: - service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} - project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} - export_default_credentials: true + credentials_json: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 - name: Build and Run Tests env: diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 69c397ef..5cfe6340 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -26,12 +26,14 @@ jobs: username: ${{ secrets.DOCKER_USER }} password: ${{ secrets.DOCKER_PASS }} - - name: Login to Google Cloud - uses: google-github-actions/setup-gcloud@v0 + - name: Authenticate to Google Cloud + uses: google-github-actions/auth@v2 with: - service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} - project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} - export_default_credentials: true + credentials_json: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 - name: Build and Run Tests env: diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 4f69a3a5..2e429055 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -28,12 +28,14 @@ jobs: username: ${{ secrets.DOCKER_USER }} password: ${{ secrets.DOCKER_PASS }} - - name: Login to Google Cloud - uses: google-github-actions/setup-gcloud@v0 + - name: Authenticate to Google Cloud + uses: google-github-actions/auth@v2 with: - service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} - project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} - export_default_credentials: true + credentials_json: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 - name: Build and Run Tests env: diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index d604b4c1..b1597095 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -30,12 +30,14 @@ jobs: username: ${{ secrets.DOCKER_USER }} password: ${{ secrets.DOCKER_PASS }} - - name: Login to Google Cloud - uses: google-github-actions/setup-gcloud@v0 + - name: Authenticate to Google Cloud + uses: google-github-actions/auth@v2 with: - service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} - project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} - export_default_credentials: true + credentials_json: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 - name: Build and Run Tests env: diff --git a/aperturedb/ParallelQuerySet.py b/aperturedb/ParallelQuerySet.py index 119379b7..57afda1c 100644 --- a/aperturedb/ParallelQuerySet.py +++ b/aperturedb/ParallelQuerySet.py @@ -8,6 +8,8 @@ from aperturedb.ParallelQuery import ParallelQuery from aperturedb.Connector import Connector +from aperturedb.Query import QueryBuilder + from aperturedb.types import Commands, Blobs logger = logging.getLogger(__name__) @@ -234,17 +236,12 @@ def constraint_filter(single_line, single_results): logger.debug( f"prev_query = {single_line[result_number]}") - if isinstance(prev_query, dict): - prev_query_cmd = [ - cmd for cmd in prev_query.keys()][0] - elif isinstance(prev_query, list) and isinstance(prev_query[0], dict) \ - and all(map(lambda k: k in known_constraint_keywords, prev_query[0].keys())) \ - and isinstance(prev_query[1], dict): - prev_query_cmd = [ - cmd for cmd in prev_query[1].keys()][0] - else: - raise Exception( - f"Contraints only implemented with with single queries; query {result_number} not single item.") + try: + prev_query_cmd = QueryBuilder.get_query_cmd( + prev_query) + except ValueError as e: + raise ValueError( + f"Constraints only implemented with single queries; query {result_number} is not a single item.") from e target_results = single_results[result_number][prev_query_cmd] target_constraints = result_constraints[result_number] diff --git a/aperturedb/Query.py b/aperturedb/Query.py index fda25e50..03923c5b 100644 --- a/aperturedb/Query.py +++ b/aperturedb/Query.py @@ -271,6 +271,27 @@ def generate_add_query( class QueryBuilder(): + @classmethod + def get_query_cmd(cls, query: Union[dict, list]) -> str: + known_constraint_keywords = ['results', 'apply'] + if isinstance(query, dict): + if len(query) != 1: + raise ValueError("Command dict must have exactly 1 key") + return list(query.keys())[0] + elif isinstance(query, list) and len(query) == 2 and isinstance(query[0], dict) \ + and query[0] and all(k in known_constraint_keywords for k in query[0].keys()) \ + and isinstance(query[1], dict): + if len(query[1]) != 1: + raise ValueError( + "Command dict in query list must have exactly 1 key") + return list(query[1].keys())[0] + else: + q_type = type(query).__name__ + q_len = len(query) if hasattr(query, '__len__') else 'N/A' + q_keys = list(query.keys()) if isinstance(query, dict) else 'N/A' + raise ValueError( + f"Cannot extract command from query structure. Expected single-command dict or [constraints_dict, command_dict]. Received: type={q_type}, len={q_len}, keys={q_keys}") + @classmethod def find_command(self, oclass: Union[str, ObjectType], params: dict) -> dict: return self.build_command(oclass, params, "Find") diff --git a/test/test_Datawizard.py b/test/test_Datawizard.py index 204b4e28..d0eab768 100644 --- a/test/test_Datawizard.py +++ b/test/test_Datawizard.py @@ -161,3 +161,38 @@ def test_all_info_preserved(self): list(filter(lambda cmd: "AddImage" in cmd, total_commands))) == 20 assert len( list(filter(lambda cmd: "FindImage" in cmd, total_commands))) == 10 + + def test_get_query_cmd_dict(self): + from aperturedb.Query import QueryBuilder + assert QueryBuilder.get_query_cmd( + {"AddImage": {"properties": {"id": 1}}}) == "AddImage" + + def test_get_query_cmd_list_with_constraints(self): + from aperturedb.Query import QueryBuilder + assert QueryBuilder.get_query_cmd( + [{"apply": True}, {"FindImage": {}}]) == "FindImage" + + def test_get_query_cmd_empty_dict(self): + from aperturedb.Query import QueryBuilder + import pytest + with pytest.raises(ValueError, match="Command dict must have exactly 1 key"): + QueryBuilder.get_query_cmd({}) + + def test_get_query_cmd_empty_list(self): + from aperturedb.Query import QueryBuilder + import pytest + with pytest.raises(ValueError, match="Cannot extract command"): + QueryBuilder.get_query_cmd([]) + + def test_get_query_cmd_list_empty_second_dict(self): + from aperturedb.Query import QueryBuilder + import pytest + with pytest.raises(ValueError, match="Command dict in query list must have exactly 1 key"): + QueryBuilder.get_query_cmd( + [{"results": {1: {"AddImage": ["==", 0]}}}, {}]) + + def test_get_query_cmd_invalid_structure(self): + from aperturedb.Query import QueryBuilder + import pytest + with pytest.raises(ValueError, match="Cannot extract command"): + QueryBuilder.get_query_cmd("invalid") From 9605e4e3b2fd91a734f4240b7bce7affd44b95a4 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sun, 24 May 2026 10:48:35 -0700 Subject: [PATCH 13/19] fix: Add a timeout of 10s to requests.get and http_client.get (#673) --- aperturedb/ImageDownloader.py | 15 ++++++--- aperturedb/Sources.py | 18 +++++++++-- aperturedb/VideoDownloader.py | 15 +++++++-- test/test_Sources.py | 60 +++++++++++++++++++++++++++++++++++ 4 files changed, 100 insertions(+), 8 deletions(-) diff --git a/aperturedb/ImageDownloader.py b/aperturedb/ImageDownloader.py index eb081136..133c0089 100644 --- a/aperturedb/ImageDownloader.py +++ b/aperturedb/ImageDownloader.py @@ -107,10 +107,13 @@ def download_image(self, url, filename): downloaded = False while True: try: - imgdata = requests.get(url) + imgdata = requests.get(url, timeout=10) downloaded = True - except requests.exceptions.ConnectionError: - logger.warning("Error with GET.", exc_info=True) + except requests.exceptions.RequestException as e: + imgdata = None + downloaded = False + logger.warning( + "Error with GET for url: %s (retry %d): %s", url, retries, e) if downloaded and imgdata is not None and imgdata.ok: break @@ -136,8 +139,12 @@ def download_image(self, url, filename): logger.exception(f"Downloaded image cannot be decoded: {url}") os.remove(filename) self.error_counter += 1 + elif not downloaded: + logger.error( + f"Failed to download URL (connection/timeout error): {url}") + self.error_counter += 1 else: - logger.error(f"URL not found: {url}") + logger.error(f"URL not found or returned error: {url}") self.error_counter += 1 self.times_arr.append(time.time() - start) diff --git a/aperturedb/Sources.py b/aperturedb/Sources.py index 803399f4..89caf449 100644 --- a/aperturedb/Sources.py +++ b/aperturedb/Sources.py @@ -45,8 +45,22 @@ def load_from_http_url(self, url, validator): retries = 0 while True: - imgdata = self.http_client.get(url) - if imgdata.ok and ("Content-Length" not in imgdata.headers or int(imgdata.headers["Content-Length"]) == imgdata.raw._fp_bytes_read): + try: + try: + imgdata = self.http_client.get(url, timeout=10) + except TypeError as e: + if "unexpected keyword argument" in str(e) and "timeout" in str(e): + imgdata = self.http_client.get(url) + else: + raise + + success = imgdata.ok and ("Content-Length" not in imgdata.headers or int( + imgdata.headers["Content-Length"]) == len(imgdata.content)) + except requests.exceptions.RequestException as e: + logger.warning(f"Error downloading object: {url} - {e}") + success = False + + if success: imgbuffer = np.frombuffer(imgdata.content, dtype='uint8') if not validator(imgbuffer): logger.error(f"VALIDATION ERROR: {url}") diff --git a/aperturedb/VideoDownloader.py b/aperturedb/VideoDownloader.py index 3b09fc86..073c40f3 100644 --- a/aperturedb/VideoDownloader.py +++ b/aperturedb/VideoDownloader.py @@ -96,7 +96,14 @@ def download_video(self, url, filename): if not os.path.exists(folder): os.makedirs(folder, exist_ok=True) - videodata = requests.get(url) + try: + videodata = requests.get(url, timeout=10) + except requests.exceptions.RequestException as e: + print(f"Error with GET for url: {url} - {e}") + self.error_counter += 1 + self.times_arr.append(time.time() - start) + return + if videodata.ok: fd = open(filename, "wb") fd.write(videodata.content) @@ -134,8 +141,12 @@ def print_stats(self): print("====== ApertureDB VideoDownloader Stats ======") times = np.array(self.times_arr) + if len(times) <= 0: + print("Error: No downloads.") + return + print("Avg Video download time(s):", np.mean(times)) - print("Img download time std:", np.std(times)) + print("Video download time std:", np.std(times)) print("Avg download throughput (videos/s)):", 1 / np.mean(times) * self.numthreads) diff --git a/test/test_Sources.py b/test/test_Sources.py index 067d1c4e..c21ef509 100644 --- a/test/test_Sources.py +++ b/test/test_Sources.py @@ -1,5 +1,6 @@ import unittest from unittest.mock import patch, MagicMock +import requests from aperturedb.Sources import Sources import botocore.exceptions from botocore import UNSIGNED @@ -162,6 +163,65 @@ def test_gs_non_auth_error_retries(self, mock_client_factory, mock_sleep): self.assertEqual(mock_blob.download_as_bytes.call_count, 2) mock_client_factory.create_anonymous_client.assert_not_called() + def test_http_passes_timeout(self): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.ok = True + mock_response.headers = {} + mock_response.content = b'mock_data' + mock_client.get.return_value = mock_response + + self.sources.http_client = mock_client + success, img = self.sources.load_from_http_url( + "http://example.com/image.jpg", self.validator) + + self.assertTrue(success) + mock_client.get.assert_called_once_with( + "http://example.com/image.jpg", timeout=10) + + def test_http_typeerror_fallback(self): + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.ok = True + mock_response.headers = {} + mock_response.content = b'mock_data' + + # Raise TypeError with expected message on first call with timeout, succeed on second without + def side_effect(*args, **kwargs): + if 'timeout' in kwargs: + raise TypeError( + "get() got an unexpected keyword argument 'timeout'") + return mock_response + + mock_client.get.side_effect = side_effect + + self.sources.http_client = mock_client + success, img = self.sources.load_from_http_url( + "http://example.com/image.jpg", self.validator) + + self.assertTrue(success) + self.assertEqual(mock_client.get.call_count, 2) + mock_client.get.assert_any_call( + "http://example.com/image.jpg", timeout=10) + mock_client.get.assert_any_call("http://example.com/image.jpg") + + @patch('time.sleep', return_value=None) + def test_http_request_exception_retries(self, mock_sleep): + mock_client = MagicMock() + + # Raise RequestException on all calls + mock_client.get.side_effect = requests.exceptions.RequestException( + "connection error") + + self.sources.http_client = mock_client + success, img = self.sources.load_from_http_url( + "http://example.com/image.jpg", self.validator) + + self.assertFalse(success) + self.assertIsNone(img) + # Should try initial + retries (1 retry = 2 calls total) + self.assertEqual(mock_client.get.call_count, 2) + class TestSourcesCaching(unittest.TestCase): @patch('boto3.client') From dfd52b6dca4e1690ce1fe54f966f84e1d49f668f Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sun, 24 May 2026 19:11:22 -0700 Subject: [PATCH 14/19] Fix: Make ApertureDBDataset applicable to Videos as well (#726) --- aperturedb/PyTorchDataset.py | 125 ++++++++++++++-------- aperturedb/TensorFlowDataset.py | 177 ++++++++++++++++++++++---------- test/test_tf_connector.py | 72 ++++++++++++- test/test_torch_connector.py | 60 ++++++++++- 4 files changed, 328 insertions(+), 106 deletions(-) diff --git a/aperturedb/PyTorchDataset.py b/aperturedb/PyTorchDataset.py index ed29fe65..b7e4fdff 100644 --- a/aperturedb/PyTorchDataset.py +++ b/aperturedb/PyTorchDataset.py @@ -15,56 +15,89 @@ class ApertureDBDataset(data.Dataset): """ This class implements a PyTorch Dataset for ApertureDB. - It is used to load images from ApertureDB into a PyTorch model. + It is used to load blobs returned by a `Find*` command from ApertureDB into a PyTorch model. It can be initialized with a query that will be used to retrieve - the images from ApertureDB. + the blobs from ApertureDB. Note that only `FindImage` blobs are decoded via OpenCV. """ - def __init__(self, client: Connector, query, label_prop=None, batch_size=1): + def __init__(self, client: Connector, query, label_prop=None, batch_size=1, command_idx=None): + + import copy self.client = client.clone() - self.query = query - self.find_image_idx = None + self.query = copy.deepcopy(query) + self.command_idx = command_idx + self.command_name = None self.total_elements = 0 - self.batch_size = batch_size - self.batch_images = [] - self.batch_start = 0 - self.batch_end = 0 - self.label_prop = label_prop - - for i in range(len(query)): - - name = list(query[i].keys())[0] - if name == "FindImage": - self.find_image_idx = i - - if self.find_image_idx is None: - logger.error( - "Query error. The query must contain one FindImage command") - raise Exception('Query Error') - - if not "results" in self.query[self.find_image_idx]["FindImage"]: - self.query[self.find_image_idx]["FindImage"]["results"] = {} + self.batch_size = batch_size + self.batch_blobs = [] + self.batch_start = 0 + self.batch_end = 0 + self.label_prop = label_prop + + allowed_find_commands = { + "FindImage", "FindVideo", "FindBlob", + "FindDescriptor", "FindBoundingBox" + } + + if self.command_idx is not None: + if not (0 <= self.command_idx < len(query)): + raise ValueError( + f"command_idx {self.command_idx} is out of range.") + self.command_name = list(query[self.command_idx].keys())[0] + if self.command_name not in allowed_find_commands: + raise ValueError( + f"Command at index {self.command_idx} is " + f"{self.command_name}, which is not a supported blob-returning Find* command.") + else: + for i in range(len(query)): + name = list(query[i].keys())[0] + if name in allowed_find_commands: + if self.command_idx is not None: + logger.warning( + "Multiple Find commands found. Selected %s at index %s.", self.command_name, self.command_idx) + break + self.command_idx = i + self.command_name = name + + if self.command_idx is None: + msg = "Query error. The query must contain at least one supported blob-returning Find command (e.g., FindImage, FindVideo, FindBlob). The first one encountered will be used." + logger.error(msg) + raise ValueError(msg) + + if "results" not in self.query[self.command_idx][self.command_name]: + self.query[self.command_idx][self.command_name]["results"] = {} if self.label_prop is not None: - results = self.query[self.find_image_idx]["FindImage"]["results"] + results = self.query[self.command_idx][self.command_name]["results"] if "list" not in results: results["list"] = [] if self.label_prop not in results["list"]: results["list"].append(self.label_prop) - self.query[self.find_image_idx]["FindImage"]["batch"] = {} - self.query[self.find_image_idx]["FindImage"]["blobs"] = True + for i in range(len(self.query)): + name = list(self.query[i].keys())[0] + if name in allowed_find_commands and i != self.command_idx: + self.query[i][name]["blobs"] = False + + self.query[self.command_idx][self.command_name]["batch"] = {} + self.query[self.command_idx][self.command_name]["blobs"] = False try: _, r, b = execute_query( client=self.client, query=self.query, blobs=[]) - batch = r[self.find_image_idx]["FindImage"]["batch"] - self.total_elements = batch["total_elements"] + resp = r[self.command_idx][self.command_name] + if resp.get("status", 0) != 0: + raise Exception( + f"Query Error: {resp.get('status')} {resp.get('info', '')}") + self.total_elements = resp.get("batch", {}).get( + "total_elements", resp.get("returned", 0)) except: logger.error( f"Query error: {self.query} {self.client.get_last_response_str()}") raise + finally: + self.query[self.command_idx][self.command_name]["blobs"] = True def __getitem__(self, index): @@ -75,14 +108,17 @@ def __getitem__(self, index): self.get_batch(index) idx = index % self.batch_size - img = self.batch_images[idx] + blob = self.batch_blobs[idx] label = self.batch_labels[idx] - nparr = np.frombuffer(img, dtype=np.uint8) - img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + if self.command_name == "FindImage": + nparr = np.frombuffer(blob, dtype=np.uint8) + blob = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + if blob is None: + raise ValueError(f"Failed to decode image at index {index}.") + blob = cv2.cvtColor(blob, cv2.COLOR_BGR2RGB) - return img, label + return blob, label def __len__(self): @@ -103,12 +139,12 @@ def get_batch(self, index): if batch_idx >= total_batches: raise Exception("Index out of range") - query = self.query - qbatch = query[self.find_image_idx]["FindImage"]["batch"] + query = self.query + qbatch = query[self.command_idx][self.command_name]["batch"] qbatch["batch_size"] = self.batch_size - qbatch["batch_id"] = batch_idx + qbatch["batch_id"] = batch_idx - query[self.find_image_idx]["FindImage"]["batch"] = qbatch + query[self.command_idx][self.command_name]["batch"] = qbatch try: @@ -128,16 +164,21 @@ def get_batch(self, index): _, r, b = execute_query( query=self.query, blobs=[], client=self.client) + resp = r[self.command_idx][self.command_name] + if resp.get("status", 0) != 0: + raise Exception( + f"Query Error: {resp.get('status')} {resp.get('info', '')}") + if len(b) == 0: logger.error(f"index: {index}") raise Exception("No results returned from ApertureDB") - self.batch_images = b - self.batch_start = self.batch_size * batch_idx - self.batch_end = self.batch_start + len(b) + self.batch_blobs = b + self.batch_start = self.batch_size * batch_idx + self.batch_end = self.batch_start + len(b) if self.label_prop: - entities = r[self.find_image_idx]["FindImage"]["entities"] + entities = r[self.command_idx][self.command_name]["entities"] try: self.batch_labels = [l[self.label_prop] for l in entities] except KeyError: diff --git a/aperturedb/TensorFlowDataset.py b/aperturedb/TensorFlowDataset.py index 12584457..7ff595c1 100644 --- a/aperturedb/TensorFlowDataset.py +++ b/aperturedb/TensorFlowDataset.py @@ -12,61 +12,90 @@ class ApertureDBTensorFlowDataset: """ This class implements a TensorFlow Dataset for ApertureDB. - It is used to load images from ApertureDB into a TensorFlow model. + It is used to load blobs returned by a `Find*` command from ApertureDB into a TensorFlow model. It can be initialized with a query that will be used to retrieve - the images from ApertureDB. + the blobs from ApertureDB. Note that only `FindImage` blobs are decoded via OpenCV. """ - def __init__(self, client: Connector, query, label_prop=None, batch_size=1): + def __init__(self, client: Connector, query, label_prop=None, batch_size=1, command_idx=None): + + import copy self.client = client.clone() - self.query = query - self.find_image_idx = None + self.query = copy.deepcopy(query) + self.command_idx = command_idx + self.command_name = None self.total_elements = 0 - self.batch_size = batch_size - self.batch_images = [] - self.batch_start = 0 - self.batch_end = 0 - self.label_prop = label_prop - self.label_type = None - - find_image_count = 0 - for i in range(len(query)): - - name = list(query[i].keys())[0] - if name == "FindImage": - self.find_image_idx = i - find_image_count += 1 - - if find_image_count != 1: - logger.error( - "Query error. The query must contain exactly one FindImage command") - raise Exception('Query Error') - - if not "results" in self.query[self.find_image_idx]["FindImage"]: - self.query[self.find_image_idx]["FindImage"]["results"] = {} + self.batch_size = batch_size + self.batch_blobs = [] + self.batch_start = 0 + self.batch_end = 0 + self.label_prop = label_prop + self.label_type = None + + allowed_find_commands = { + "FindImage", "FindVideo", "FindBlob", + "FindDescriptor", "FindBoundingBox" + } + + if self.command_idx is not None: + if not (0 <= self.command_idx < len(query)): + raise ValueError( + f"command_idx {self.command_idx} is out of range.") + self.command_name = list(query[self.command_idx].keys())[0] + if self.command_name not in allowed_find_commands: + raise ValueError( + f"Command at index {self.command_idx} is " + f"{self.command_name}, which is not a supported blob-returning Find* command.") + else: + for i in range(len(query)): + name = list(query[i].keys())[0] + if name in allowed_find_commands: + if self.command_idx is not None: + logger.warning( + "Multiple Find commands found. Selected %s at index %s.", self.command_name, self.command_idx) + break + self.command_idx = i + self.command_name = name + + if self.command_idx is None: + msg = "Query error. The query must contain at least one supported blob-returning Find command (e.g., FindImage, FindVideo, FindBlob). The first one encountered will be used." + logger.error(msg) + raise ValueError(msg) + + if "results" not in self.query[self.command_idx][self.command_name]: + self.query[self.command_idx][self.command_name]["results"] = {} if self.label_prop is not None: - results = self.query[self.find_image_idx]["FindImage"]["results"] + results = self.query[self.command_idx][self.command_name]["results"] if "list" not in results: results["list"] = [] if self.label_prop not in results["list"]: results["list"].append(self.label_prop) - self.query[self.find_image_idx]["FindImage"]["batch"] = {} - self.query[self.find_image_idx]["FindImage"]["blobs"] = False + for i in range(len(self.query)): + name = list(self.query[i].keys())[0] + if name in allowed_find_commands and i != self.command_idx: + self.query[i][name]["blobs"] = False + + self.query[self.command_idx][self.command_name]["batch"] = {} + self.query[self.command_idx][self.command_name]["blobs"] = False try: _, r, b = execute_query( client=self.client, query=self.query, blobs=[]) - batch = r[self.find_image_idx]["FindImage"]["batch"] - self.total_elements = batch["total_elements"] + resp = r[self.command_idx][self.command_name] + if resp.get("status", 0) != 0: + raise Exception( + f"Query Error: {resp.get('status')} {resp.get('info', '')}") + self.total_elements = resp.get("batch", {}).get( + "total_elements", resp.get("returned", 0)) except: logger.error( f"Query error: {self.query} {self.client.get_last_response_str()}") raise - - self.query[self.find_image_idx]["FindImage"]["blobs"] = True + finally: + self.query[self.command_idx][self.command_name]["blobs"] = True def is_in_range(self, index): @@ -83,12 +112,12 @@ def get_batch(self, index): if batch_idx >= total_batches: raise Exception("Index out of range") - query = self.query - qbatch = query[self.find_image_idx]["FindImage"]["batch"] + query = self.query + qbatch = query[self.command_idx][self.command_name].get("batch", {}) qbatch["batch_size"] = self.batch_size - qbatch["batch_id"] = batch_idx + qbatch["batch_id"] = batch_idx - query[self.find_image_idx]["FindImage"]["batch"] = qbatch + query[self.command_idx][self.command_name]["batch"] = qbatch try: @@ -108,16 +137,21 @@ def get_batch(self, index): _, r, b = execute_query( query=self.query, blobs=[], client=self.client) + resp = r[self.command_idx][self.command_name] + if resp.get("status", 0) != 0: + raise Exception( + f"Query Error: {resp.get('status')} {resp.get('info', '')}") + if len(b) == 0: logger.error(f"index: {index}") raise Exception("No results returned from ApertureDB") - self.batch_images = b - self.batch_start = self.batch_size * batch_idx - self.batch_end = self.batch_start + len(b) + self.batch_blobs = b + self.batch_start = self.batch_size * batch_idx + self.batch_end = self.batch_start + len(b) if self.label_prop: - entities = r[self.find_image_idx]["FindImage"]["entities"] + entities = r[self.command_idx][self.command_name]["entities"] try: self.batch_labels = [l[self.label_prop] for l in entities] except KeyError: @@ -140,34 +174,69 @@ def generator(self): self.get_batch(index) idx = index % self.batch_size - img = self.batch_images[idx] + blob = self.batch_blobs[idx] label = self.batch_labels[idx] - nparr = np.frombuffer(img, dtype=np.uint8) - img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) - img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + if self.command_name == "FindImage": + nparr = np.frombuffer(blob, dtype=np.uint8) + blob = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + if blob is None: + raise ValueError( + f"Failed to decode image at index {index}.") + blob = cv2.cvtColor(blob, cv2.COLOR_BGR2RGB) - yield img, label + yield blob, label def get_dataset(self): import tensorflow as tf if self.label_type is None: if self.total_elements > 0: - self.get_batch(0) - if isinstance(self.batch_labels[0], int): - self.label_type = tf.int32 - elif isinstance(self.batch_labels[0], float): - self.label_type = tf.float32 - else: + # Infer label_type with a lightweight query (blobs=False, batch_size=1) + import copy + infer_query = copy.deepcopy(self.query) + infer_query[self.command_idx][self.command_name]["blobs"] = False + infer_query[self.command_idx][self.command_name].setdefault( + "batch", {}) + infer_query[self.command_idx][self.command_name]["batch"]["batch_size"] = 1 + infer_query[self.command_idx][self.command_name]["batch"]["batch_id"] = 0 + + try: + _, r, _ = execute_query( + query=infer_query, blobs=[], client=self.client) + resp = r[self.command_idx][self.command_name] + if resp.get("status", 0) != 0: + raise Exception( + f"Query Error: {resp.get('status')} {resp.get('info', '')}") + if self.label_prop and "entities" in resp and len(resp["entities"]) > 0: + sample_label = resp["entities"][0].get(self.label_prop) + else: + sample_label = "none" + + if isinstance(sample_label, int): + self.label_type = tf.int32 + elif isinstance(sample_label, float): + self.label_type = tf.float32 + else: + self.label_type = tf.string + except Exception as e: + logger.warning( + "Failed to infer label_type: %s. Defaulting to tf.string.", e) self.label_type = tf.string else: self.label_type = tf.string + if self.command_name == "FindImage": + tensor_shape = (None, None, 3) + tensor_dtype = tf.uint8 + else: + tensor_shape = () + tensor_dtype = tf.string + return tf.data.Dataset.from_generator( self.generator, output_signature=( - tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), + tf.TensorSpec(shape=tensor_shape, dtype=tensor_dtype), tf.TensorSpec(shape=(), dtype=self.label_type) ) ) diff --git a/test/test_tf_connector.py b/test/test_tf_connector.py index f7ac96c2..ecbe803d 100644 --- a/test/test_tf_connector.py +++ b/test/test_tf_connector.py @@ -15,9 +15,13 @@ def validate_dataset(self, dataset: tf.data.Dataset, expected_length): count = 0 # Iterate over dataset. - for img, label in dataset: - if tf.size(img).numpy() == 0: - logger.error("Empty image?") + for data, label in dataset: + if data.dtype == tf.string: + size = tf.strings.length(data).numpy() + else: + size = tf.size(data).numpy() + if size == 0: + logger.error("Empty data?") assert False count += 1 assert count == expected_length @@ -85,8 +89,8 @@ def test_batchedDataset(self, db, utils, images): start = time.time() count = 0 - for imgs, labels in batched_dataset: - count += tf.shape(imgs)[0].numpy() + for data, labels in batched_dataset: + count += tf.shape(data)[0].numpy() assert count == len_limit time_taken = time.time() - start @@ -137,3 +141,61 @@ def side_effect_float(*args, **kwargs): ), [{"FindImage": {"results": {"list": ["prop"]}}}], label_prop="prop") dataset_f = dataset_wrapper_f.get_dataset() assert dataset_f.element_spec[1].dtype == tf.float32 + + def test_findBlob(self, db, utils, insert_data_from_csv): + blobs, _ = insert_data_from_csv("./input/blobs.adb.csv") + assert len(blobs) > 0 + query = [{ + "FindBlob": { + "results": {} + } + }] + + dataset_wrapper = ApertureDBTensorFlowDataset( + db, query, label_prop="license") + dataset = dataset_wrapper.get_dataset() + + count = 0 + for blob, label in dataset: + assert isinstance(blob.numpy(), bytes) + assert label.dtype == tf.int32 + count += 1 + + assert count == utils.count_entities("_Blob") + + def test_findVideo_mocked(self): + from unittest.mock import patch + import tensorflow as tf + + class DummyClient: + def clone(self): + return self + + def get_last_response_str(self): + return "" + + query = [{"FindVideo": {"results": {"list": ["prop"]}}}] + + with patch('aperturedb.TensorFlowDataset.execute_query') as mock_exec: + def side_effect(*args, **kwargs): + batch_dict = {"total_elements": 1} + entities = [{"prop": 1}] + r = [{"FindVideo": {"batch": batch_dict, "entities": entities}}] + b = [b"mock_video_bytes"] + return None, r, b + + mock_exec.side_effect = side_effect + dataset_wrapper = ApertureDBTensorFlowDataset( + DummyClient(), query, label_prop="prop") + dataset = dataset_wrapper.get_dataset() + + assert dataset.element_spec[1].dtype == tf.int32 + assert dataset.element_spec[0].dtype == tf.string + + count = 0 + for data, label in dataset: + assert isinstance(data.numpy(), bytes) + assert data.numpy() == b"mock_video_bytes" + assert label.numpy() == 1 + count += 1 + assert count == 1 diff --git a/test/test_torch_connector.py b/test/test_torch_connector.py index 22dd91af..f33d64a3 100644 --- a/test/test_torch_connector.py +++ b/test/test_torch_connector.py @@ -20,11 +20,11 @@ def validate_dataset(self, dataset: Union[DataLoader, Dataset], expected_length) count = 0 # Iterate over dataset. - for img in dataset: - if len(img[0]) < 0: - logger.error("Empty image?") - assert True == False - count += len(img[1]) if isinstance(dataset, DataLoader) else 1 + for data in dataset: + if len(data[0]) == 0: + logger.error("Empty data?") + assert False + count += len(data[1]) if isinstance(dataset, DataLoader) else 1 assert count == expected_length time_taken = time.time() - start @@ -58,6 +58,25 @@ def test_nativeConstraints(self, db, utils, images): self.validate_dataset(dataset, utils.count_images()) + def test_findBlob(self, db, utils, insert_data_from_csv): + blobs, _ = insert_data_from_csv("./input/blobs.adb.csv") + assert len(blobs) > 0 + query = [{ + "FindBlob": { + "results": {} + } + }] + + dataset = PyTorchDataset.ApertureDBDataset( + db, query, label_prop="license") + + assert len(dataset) == utils.count_entities("_Blob") + for blob, label in dataset: + # For FindBlob, the return is raw bytes and label should be an int when label_prop='license' + assert isinstance(blob, bytes) + assert isinstance(label, int) + break + def test_datasetWithMultiprocessing(self, db, utils, images): len_limit = utils.count_images() # This is a hack against a bug in batch API. @@ -124,3 +143,34 @@ def test_datasetWithMultiprocessing(self, db, utils, images): self.validate_dataset(data_loader, len_limit) dist.destroy_process_group() + + def test_findVideo_mocked(self): + from unittest.mock import patch + + class DummyClient: + def clone(self): + return self + + def get_last_response_str(self): + return "" + + query = [{"FindVideo": {"results": {"list": ["prop"]}}}] + + with patch('aperturedb.PyTorchDataset.execute_query') as mock_exec: + def side_effect(*args, **kwargs): + batch_dict = {"total_elements": 1} + entities = [{"prop": 1}] + r = [{"FindVideo": {"batch": batch_dict, "entities": entities}}] + b = [b"mock_video_bytes"] + return None, r, b + + mock_exec.side_effect = side_effect + dataset = PyTorchDataset.ApertureDBDataset( + DummyClient(), query, label_prop="prop") + + assert len(dataset) == 1 + for blob, label in dataset: + assert isinstance(blob, bytes) + assert blob == b"mock_video_bytes" + assert label == 1 + break From 9a096a06b17c901f84697b235eaa7549bcb064c4 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sun, 24 May 2026 19:12:41 -0700 Subject: [PATCH 15/19] feat: enhance error logging and add error_handler to queries (#727) --- aperturedb/CommonLibrary.py | 37 +++++++++--- aperturedb/ParallelQuery.py | 6 +- aperturedb/ParallelQuerySet.py | 28 ++++++--- aperturedb/VideoDataCSV.py | 5 +- test/test_ResponseHandler.py | 101 +++++++++++++++++++++++++++++++++ 5 files changed, 155 insertions(+), 22 deletions(-) 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 57afda1c..529a0999 100644 --- a/aperturedb/ParallelQuerySet.py +++ b/aperturedb/ParallelQuerySet.py @@ -48,23 +48,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) @@ -270,7 +279,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 @@ -300,7 +309,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 From d5c4f00b8de643fccc0fe11591bd6f14d1d940bf Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Sun, 24 May 2026 19:14:59 -0700 Subject: [PATCH 16/19] fix: default unbound variables in run_test.sh (#715) Fixes issue with unbound variables like GCP_SERVICE_ACCOUNT_KEY when running test/run_test.sh locally with set -u. * fix: ensure GOOGLE_APPLICATION_CREDENTIALS is set to avoid metadata server hang, and simplify pytest args * docs: update robots.md testing guide for native skip support --- robots.md | 16 +++++----------- test/run_test.sh | 45 +++++++++++++++++++++++++++++++++------------ 2 files changed, 38 insertions(+), 23 deletions(-) diff --git a/robots.md b/robots.md index 966408db..9c36282c 100644 --- a/robots.md +++ b/robots.md @@ -55,21 +55,15 @@ Expected results with a clean environment: ~118 passed + 4 skipped for non-HTTP, A subset of the `pytest` tests interact with external objects, such as fetching blob configurations directly from cloud providers (AWS S3, Google Storage - GS). If your local environment or the automated agent lacks valid AWS/GCP credentials (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `GCP_SERVICE_ACCOUNT_KEY`), you will hit authentication validation errors when resolving `storage.Client()` architectures or boto3 structures during execution. -To dynamically skip external storage tests, modify the `pytest` invocation inside `test/run_test.sh` before running. Because `run_test_container.sh` mounts this file directly into the test container at runtime, no image rebuild is needed — just edit the file and re-run. - -Use this quick `sed` command to patch the exclusion: +To skip these tests, set `SKIP_SLOW_TESTS=true` before running `run_test_container.sh`. This environment variable is passed into the container and automatically appends `and not slow and not external_network` to the `pytest` marker expression. ```bash cd test -# Modify the pytest invocation to ignore AWS/GCP specific loaders and remotely constrained mark configurations -sed -i 's/pytest --cov=aperturedb -m "\$FILTER"/pytest --cov=aperturedb -k "not test_S3ImageLoader and not test_GSImageLoader and not test_S3VideoLoader and not test_GSVideoLoader" -m "\$FILTER and not remote_credentials and not external_network"/' run_test.sh - -# Then run the tests as usual set -a && source .env && set +a -bash run_test_container.sh +SKIP_SLOW_TESTS=true bash run_test_container.sh ``` -By applying this change, any LLM/agent can successfully execute internal `aperturedb` connector + python-centric API tests without raising failures tied strictly to external cloud requirements. +By setting this variable, any LLM/agent can successfully execute internal `aperturedb` connector + python-centric API tests without raising failures tied strictly to external cloud requirements. --- @@ -92,7 +86,7 @@ These warnings appear during a successful local run and can be ignored: | **DB image** | `aperturedata/aperturedb-community:latest` | `aperturedata/aperturedb:dev` | | **Lenz tag** | `latest` | `dev` | | **Test filtering** | Skips `remote_credentials`, `external_network`, S3/GCS loaders | Runs all tests including cloud-dependent ones | -| **AWS/GCP credentials** | Not present — must patch `run_test.sh` | Present as GitHub secrets | +| **AWS/GCP credentials** | Not present — use `SKIP_SLOW_TESTS=true` | Present as GitHub secrets | | **Log upload on failure** | Skipped (no AWS creds) | Uploads to S3 bucket `python-ci-runs` | | **`BUILD_AUX_IMAGES`** | `true` (default in `ci.sh`) — builds notebook/coverage images | `false` — skips aux images to save time | | **`TEST_PROTOCOL`** | `both` (default in `run_test_container.sh`) | `both` (explicit in `pr.yaml`) | @@ -102,4 +96,4 @@ These warnings appear during a successful local run and can be ignored: | **Kaggle credentials** | Dummy (`KAGGLE_username=ci`, `KAGGLE_key=dummy`) | Same dummy values | | **Coverage HTML** | `coverage html` fails (no data in container path) | Same issue; coverage not collected meaningfully | | **Concurrency** | N/A | New push to PR cancels in-progress run | -| **`run_test.sh` override** | Mounted as volume — edit `test/run_test.sh` locally (e.g. apply credentials filter from Section 3), no image rebuild needed | Also mounted as volume — the committed `test/run_test.sh` runs in CI. Do NOT commit the credentials filter; CI needs full test coverage | +| **`run_test.sh` override** | Mounted as volume — edits locally apply immediately | Also mounted as volume | diff --git a/test/run_test.sh b/test/run_test.sh index 58d68a15..1cc9633c 100755 --- a/test/run_test.sh +++ b/test/run_test.sh @@ -4,6 +4,10 @@ set -u set -e set -o pipefail +: "${GCP_SERVICE_ACCOUNT_KEY:=""}" +: "${APERTUREDB_LOG_PATH:=""}" +: "${FILTER:=""}" + mkdir -p output rm -rf output/* mkdir -p input/blobs @@ -22,25 +26,42 @@ python3 generateInput.py echo "Done generating input files." echo "Running tests..." -CREDENTIALS_FILE='/tmp/key.json' -echo $GCP_SERVICE_ACCOUNT_KEY > $CREDENTIALS_FILE -export GOOGLE_APPLICATION_CREDENTIALS=$CREDENTIALS_FILE +if [ -n "$GCP_SERVICE_ACCOUNT_KEY" ]; then + CREDENTIALS_FILE=$(mktemp "${TMPDIR:-/tmp}/key.XXXXXX") + trap 'rm -f "$CREDENTIALS_FILE"' EXIT + printf "%s\n" "$GCP_SERVICE_ACCOUNT_KEY" > "$CREDENTIALS_FILE" + export GOOGLE_APPLICATION_CREDENTIALS="$CREDENTIALS_FILE" +else + unset GOOGLE_APPLICATION_CREDENTIALS +fi # capture errors set +e -CLIENT_PATH="${APERTUREDB_LOG_PATH}/../client/${FILTER}" -CLIENT_PATH=${CLIENT_PATH// /_} -mkdir -p ${CLIENT_PATH} +SAFE_FILTER=$(printf "%s" "$FILTER" | tr -c 'a-zA-Z0-9_-' '_') +if [ -n "$APERTUREDB_LOG_PATH" ]; then + CLIENT_PATH="${APERTUREDB_LOG_PATH}/../client/${SAFE_FILTER}" +else + CLIENT_PATH="output/client/${SAFE_FILTER}" +fi +mkdir -p "$CLIENT_PATH" if [ "${SKIP_SLOW_TESTS:-false}" == "true" ]; then echo "Skipping slow and external tests for this run..." - pytest_filter="${FILTER} and not slow and not external_network" + if [ -n "$FILTER" ]; then + pytest_filter="${FILTER} and not slow and not external_network" + else + pytest_filter="not slow and not external_network" + fi else pytest_filter="${FILTER}" fi -PROJECT=aperturedata KAGGLE_username=ci KAGGLE_key=dummy python3 -m pytest --cov=aperturedb -m "$pytest_filter" test_*.py -v | tee ${CLIENT_PATH}/test.log +if [ -n "$pytest_filter" ]; then + PROJECT=aperturedata KAGGLE_username=ci KAGGLE_key=dummy python3 -m pytest --cov=aperturedb -m "$pytest_filter" test_*.py -v | tee "${CLIENT_PATH}/test.log" +else + PROJECT=aperturedata KAGGLE_username=ci KAGGLE_key=dummy python3 -m pytest --cov=aperturedb test_*.py -v | tee "${CLIENT_PATH}/test.log" +fi RESULT=$? -cp error*.log -v ${CLIENT_PATH} || true +cp error*.log -v "$CLIENT_PATH" || true if [[ $RESULT != 0 ]]; then echo "Test failed; outputting db log:" @@ -49,9 +70,9 @@ if [[ $RESULT != 0 ]]; then BUCKET=python-ci-runs NOW=$(date -Iseconds) ARCHIVE_NAME=logs.tar.gz - DESTINATION="s3://${BUCKET}/aperturedb-${NOW}-${FILTER// /_}.tgz" - tar czf ${ARCHIVE_NAME} ${APERTUREDB_LOG_PATH}/.. - aws s3 cp ${ARCHIVE_NAME} $DESTINATION + DESTINATION="s3://${BUCKET}/aperturedb-${NOW}-${SAFE_FILTER}.tgz" + tar czf "${ARCHIVE_NAME}" "${APERTUREDB_LOG_PATH}/.." + aws s3 cp "${ARCHIVE_NAME}" "$DESTINATION" echo "Log output to $DESTINATION" else echo "Unable to output log, APERTUREDB_LOG_PATH not set." From 13c803cd30aa5dc6993ff33757a43e8486325c64 Mon Sep 17 00:00:00 2001 From: ad-claw000 Date: Tue, 26 May 2026 06:21:13 -0700 Subject: [PATCH 17/19] fix: add explicit close() method to connectors (#736) --- aperturedb/Connector.py | 12 ++++--- aperturedb/ConnectorRest.py | 29 +++++++++++----- aperturedb/ParallelQuery.py | 66 +++++++++++++++++++++--------------- test/test_ConnectionPool.py | 25 ++++++++++++++ test/test_Connector.py | 13 +++++++ test/test_Parallel.py | 49 ++++++++++++++++++++++++++ test/test_UserConvenience.py | 47 +++++++++++++++++++++++++ 7 files changed, 202 insertions(+), 39 deletions(-) diff --git a/aperturedb/Connector.py b/aperturedb/Connector.py index 1bd71b42..50285973 100644 --- a/aperturedb/Connector.py +++ b/aperturedb/Connector.py @@ -114,7 +114,7 @@ class Connector(object): :::note - The connection is established only when a query is run. - - A new connection is established for each instance that runs a query, and gets closed only at destruction. + - A new connection is established for each instance that runs a query, and gets closed explicitly by calling `close()` or at destruction. ::: Args: @@ -228,10 +228,14 @@ def authenticate(self, shared_data, user, password, token): self.shared_data = shared_data self.authenticated = True - def __del__(self): - if self.connected: + def close(self): + if getattr(self, 'conn', None) is not None: self.conn.close() - self.connected = False + self.conn = None + self.connected = False + + def __del__(self): + self.close() def _send_msg(self, data): if len(data) > (DEFAULT_MAX_MESSAGE_SIZE_MB * 2**20): diff --git a/aperturedb/ConnectorRest.py b/aperturedb/ConnectorRest.py index c3231234..16381c75 100644 --- a/aperturedb/ConnectorRest.py +++ b/aperturedb/ConnectorRest.py @@ -117,6 +117,16 @@ def __init__(self, host="localhost", port=None, # Session is useful because it does not add "Connection: close header" # Since we will be making same call to the same URL, making a session # REF: https://requests.readthedocs.io/en/latest/user/advanced/ + self.http_session = None + self._init_session() + + self.last_response = '' + self.last_query_time = 0 + + self.url = ('https' if self.use_ssl else 'http') + \ + '://' + self.host + ':' + str(self.port) + '/api/' + + def _init_session(self): self.http_session = requests.Session() if self.config.verify_hostname: if self.config.ca_cert: @@ -125,17 +135,20 @@ def __init__(self, host="localhost", port=None, adapter = CustomHTTPAdapter(ca_cert=None) self.http_session.mount('https://', adapter=adapter) - self.last_response = '' - self.last_query_time = 0 - - self.url = ('https' if self.use_ssl else 'http') + \ - '://' + self.host + ':' + str(self.port) + '/api/' + def close(self): + if hasattr(self, 'http_session') and self.http_session is not None: + self.http_session.close() + self.http_session = None + super().close() def __del__(self): - logger.info("Done with connector REST.") - self.http_session.close() + self.close() - def _query(self, query, blob_array = [], try_resume=True): + def _query(self, query, blob_array=None, try_resume=True): + if blob_array is None: + blob_array = [] + if getattr(self, 'http_session', None) is None: + self._init_session() response_blob_array = [] # Check the query type if not isinstance(query, str): # assumes json diff --git a/aperturedb/ParallelQuery.py b/aperturedb/ParallelQuery.py index e2237420..cd3a2180 100644 --- a/aperturedb/ParallelQuery.py +++ b/aperturedb/ParallelQuery.py @@ -4,6 +4,7 @@ import numpy as np import logging import inspect +import threading from aperturedb.DaskManager import DaskManager @@ -65,6 +66,7 @@ def __init__(self, client: Connector, dry_run: bool = False): self.blobs_per_query = 0 self.daskManager = None self.batch_command = execute_query + self.error_counter_lock = threading.Lock() def generate_batch(self, data: List[Tuple[Commands, Blobs]]) -> Tuple[Commands, Blobs]: """ @@ -198,7 +200,8 @@ def response_handler(query, qblobs, resp, rblobs, qindex): return indexless_hand worker_stats["objects_existed"] = sum( v['status'] == 2 for i in r for k, v in i.items()) elif result == 1: - self.error_counter += 1 + with self.error_counter_lock: + self.error_counter += 1 worker_stats["succeeded_queries"] = 0 worker_stats["succeeded_commands"] = 0 worker_stats["objects_existed"] = 0 @@ -238,31 +241,39 @@ def worker(self, thid: int, generator, start: int, end: int, run_event) -> None: # A new connection will be created for each thread client = self.client.clone() - total_batches = (end - start) // self.batchsize - - if (end - start) % self.batchsize > 0: - total_batches += 1 - - logger.info( - f"Worker {thid} executing {total_batches} batches, {self.stats=}") - for i in range(total_batches): - if not run_event.is_set(): - break - batch_start = start + i * self.batchsize - batch_end = min(batch_start + self.batchsize, end) + try: + total_batches = (end - start) // self.batchsize - try: - self.do_batch(client, batch_start, - generator[batch_start:batch_end]) - except Exception as e: - logger.exception(e) - logger.warning( - f"Worker {thid} failed to execute batch {i}: [{batch_start},{batch_end}]") - self.error_counter += 1 + if (end - start) % self.batchsize > 0: + total_batches += 1 - if self.stats: - self.pb.update(batch_end - batch_start) - logger.info(f"Worker {thid} executed {total_batches} batches") + logger.info( + f"Worker {thid} executing {total_batches} batches, {self.stats=}") + executed_batches = 0 + for i in range(total_batches): + if not run_event.is_set(): + break + batch_start = start + i * self.batchsize + batch_end = min(batch_start + self.batchsize, end) + + try: + self.do_batch(client, batch_start, + generator[batch_start:batch_end]) + except Exception as e: + logger.exception(e) + logger.warning( + f"Worker {thid} failed to execute batch {i}: [{batch_start},{batch_end}]") + with self.error_counter_lock: + self.error_counter += 1 + + executed_batches += 1 + if self.stats: + self.pb.update(batch_end - batch_start) + logger.info(f"Worker {thid} executed {executed_batches} batches") + finally: + # Explicitly close the connection to avoid exhausting server connection limits + if client is not self.client and hasattr(client, 'close') and callable(client.close): + client.close() def get_objects_existed(self) -> int: return sum(stat["objects_existed"] @@ -290,19 +301,20 @@ def query(self, generator, batchsize: int = 1, numthreads: int = 4, stats: bool use_dask = hasattr(generator, "use_dask") and generator.use_dask if use_dask: self._reset(batchsize=batchsize, numthreads=numthreads) - self.daskmanager = DaskManager(num_workers=numthreads) + self.daskManager = DaskManager(num_workers=numthreads) if hasattr(self, "query_setup"): self.query_setup(generator) if use_dask: - results, self.total_actions_time = self.daskmanager.run( + results, self.total_actions_time = self.daskManager.run( self.__class__, self.client, generator, batchsize, stats=stats, dry_run=self.dry_run) self.actual_stats = [] for result in results: if result is not None: self.times_arr.extend(result.times_arr) - self.error_counter += result.error_counter + with self.error_counter_lock: + self.error_counter += result.error_counter self.actual_stats.append( {"succeeded_queries": result.succeeded_queries, "succeeded_commands": result.succeeded_commands, diff --git a/test/test_ConnectionPool.py b/test/test_ConnectionPool.py index 3f42094b..0794fd8e 100644 --- a/test/test_ConnectionPool.py +++ b/test/test_ConnectionPool.py @@ -81,6 +81,31 @@ def test_pool_timeout(self): with pool.get_connection(timeout=0.1): pass + def test_pool_close(self): + pool = ConnectionPool(pool_size=2, connection_factory=_make_connector) + + # Borrow one connection to ensure it connects + with pool.get_connection() as conn: + conn.query([{"GetStatus": {}}]) + # Connection is established + self.assertTrue(conn.connected) + self.assertIsNotNone(conn.conn) + borrowed_conn = conn + + # Close the pool + pool.close() + + # Verify connections are closed and in a predictable state + self.assertFalse(borrowed_conn.connected) + self.assertIsNone(borrowed_conn.conn) + self.assertEqual(pool.available(), 0) + + # Verify that after closing, the connector can still query (reconnects) + response, _ = borrowed_conn.query([{"GetStatus": {}}]) + self.assertTrue(isinstance(response, list)) + self.assertTrue(borrowed_conn.connected) + self.assertIsNotNone(borrowed_conn.conn) + if __name__ == '__main__': unittest.main() diff --git a/test/test_Connector.py b/test/test_Connector.py index 0d5cd45b..736c3733 100644 --- a/test/test_Connector.py +++ b/test/test_Connector.py @@ -29,3 +29,16 @@ def test_check_status_mixed(self): # Test case 6: multiple keys in a single dict (tests values traversal) res6 = {"FindImage": {"status": 0}, "FindEntity": {"status": -4}} assert connector.check_status(res6) == -4 + + def test_close_partially_initialized(self): + # Test that close() does not throw if called before self.conn is set. + # This simulates a failure in __init__ followed by __del__ calling close(). + class UninitializedConnector(Connector): + def __init__(self): + # Deliberately avoid calling super().__init__() so self.conn is never set + pass + + connector = UninitializedConnector() + # Should not raise AttributeError + connector.close() + assert getattr(connector, "connected", True) is False diff --git a/test/test_Parallel.py b/test/test_Parallel.py index 85c7cd51..488cc3a0 100644 --- a/test/test_Parallel.py +++ b/test/test_Parallel.py @@ -114,6 +114,55 @@ def test_dictResponseHandling(self): print(e) raise + def test_parallel_query_worker_closes_connection(self, db, monkeypatch): + from aperturedb.QueryGenerator import QueryGenerator + import threading + + class MockQueryGenerator(QueryGenerator): + def __len__(self): + return 2 + + def getitem(self, idx): + return [{"FindImage": {}}], [] + + pq = ParallelQuery(db) + + closed_count = [0] + original_clone = pq.client.clone + + def mock_clone(): + cloned = original_clone() + original_close = cloned.close + + def mock_close(): + closed_count[0] += 1 + original_close() + cloned.close = mock_close + return cloned + monkeypatch.setattr(pq.client, "clone", mock_clone) + + original_do_batch = pq.do_batch + + def mock_do_batch(client, batch_start, data): + if batch_start == 1: + raise Exception("Simulated do_batch exception") + original_do_batch(client, batch_start, data) + monkeypatch.setattr(pq, "do_batch", mock_do_batch) + + # Test exception in do_batch + pq.query(MockQueryGenerator(), batchsize=1, numthreads=1) + # Should close even if exception occurred in do_batch + assert closed_count[0] == 1 + + # Test early exit when run_event is cleared + closed_count[0] = 0 + run_event = threading.Event() + run_event.clear() # Not set, so worker breaks immediately + + # worker signature: worker(self, thid: int, generator, start: int, end: int, run_event) + pq.worker(0, MockQueryGenerator(), 0, 1, run_event) + assert closed_count[0] == 1 + def test_dask_dry_run(db: Connector): from aperturedb.ParallelLoader import ParallelLoader diff --git a/test/test_UserConvenience.py b/test/test_UserConvenience.py index 848bae2d..6c2291de 100644 --- a/test/test_UserConvenience.py +++ b/test/test_UserConvenience.py @@ -43,3 +43,50 @@ def mock_post(self, url, headers, files, verify): # Ensure that the mock post was called, 1 time to authenticate, 1 time to query assert posts == 2 Session.post = old_post + + def test_ConnectorRest_close_and_recreate_session(self): + """ + Test that ConnectorRest can be closed, clearing its http_session, + and that a subsequent call to query() will transparently recreate it. + """ + client = ConnectorRest(host="dummy", user="admin", password="password") + posts = 0 + + def mock_post(self, url, headers, files, verify): + nonlocal posts + response1 = { + "json": [{"Authenticate": { + "status": 0, + "session_token": "x", + "refresh_token": "2", + "session_token_expires_in": 3600, + "refresh_token_expires_in": 3600 + }}], + "blobs": [] + } + + r = SimpleNamespace(status_code=200, text=json.dumps(response1)) + posts += 1 + return r + + old_post = Session.post + Session.post = mock_post + + try: + # Query 1: Initialize normally + client.query("[{\"FindEntity\": {\"_ref\": 1}}]") + assert posts == 2 # 1 auth, 1 query + assert client.http_session is not None + old_session = client.http_session + + # Explicitly close the connector + client.close() + assert getattr(client, "http_session", None) is None + + # Query 2: Should transparently recreate the session + client.query("[{\"FindEntity\": {\"_ref\": 2}}]") + assert posts == 3 + assert client.http_session is not None + assert client.http_session is not old_session + finally: + Session.post = old_post From 27e6f5313df42402534b49badd4f82c6ac7b9c94 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 5 Jun 2026 02:43:14 +0000 Subject: [PATCH 18/19] Version bump: 0.4.59 to 0.4.60 --- aperturedb/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aperturedb/__init__.py b/aperturedb/__init__.py index 75adf709..a4fccee9 100644 --- a/aperturedb/__init__.py +++ b/aperturedb/__init__.py @@ -10,7 +10,7 @@ import signal import sys -__version__ = "0.4.59" +__version__ = "0.4.60" logger = logging.getLogger(__name__) From d13c10ae7807dc36395f8e3677181ec32e4cbd1e Mon Sep 17 00:00:00 2001 From: luisremis Date: Mon, 8 Jun 2026 12:30:39 -0700 Subject: [PATCH 19/19] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .github/workflows/nightly.yml | 11 ++++++----- .github/workflows/pr.yaml | 11 ++++++----- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index bb32ee55..25401e54 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -23,13 +23,14 @@ jobs: username: ${{ secrets.DOCKER_USER }} password: ${{ secrets.DOCKER_PASS }} - - name: Login to Google Cloud - uses: google-github-actions/setup-gcloud@v0 + - name: Authenticate to Google Cloud + uses: google-github-actions/auth@v2 with: - service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} - project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} - export_default_credentials: true + credentials_json: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 - name: Build and Run Tests env: # Run both protocols inside a single job. run_test_container.sh starts diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2e429055..57bea7a9 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -91,13 +91,14 @@ jobs: username: ${{ secrets.DOCKER_USER }} password: ${{ secrets.DOCKER_PASS }} - - name: Login to Google Cloud - uses: google-github-actions/setup-gcloud@v0 + - name: Authenticate to Google Cloud + uses: google-github-actions/auth@v2 with: - service_account_key: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} - project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} - export_default_credentials: true + credentials_json: ${{ secrets.GCP_SERVICE_ACCOUNT_KEY }} + project_id: ${{ secrets.GCP_SERVICE_ACCOUNT_PROJECT_ID }} + - name: Set up Cloud SDK + uses: google-github-actions/setup-gcloud@v2 - name: Build tests on pytorch GPU image run: | rm -rf docker/pytorch-gpu/aperturedata