From c647a53b7d4b6925bd5b7019931ccbe857531cb4 Mon Sep 17 00:00:00 2001 From: PatersonProjects Date: Fri, 4 Sep 2026 11:07:58 -0700 Subject: [PATCH 1/4] Add aerospike integration tests --- .../persisters/test_b_aerospike.py | 371 ++++++++++++++++++ 1 file changed, 371 insertions(+) create mode 100644 tests/integrations/persisters/test_b_aerospike.py diff --git a/tests/integrations/persisters/test_b_aerospike.py b/tests/integrations/persisters/test_b_aerospike.py new file mode 100644 index 000000000..d226e7a57 --- /dev/null +++ b/tests/integrations/persisters/test_b_aerospike.py @@ -0,0 +1,371 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import os +import pickle +import tomllib +import uuid +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + +if os.environ.get("BURR_CI_INTEGRATION_TESTS") != "true": + pytest.skip("Skipping integration tests", allow_module_level=True) + +from burr.core import state +from burr.core.persistence import BaseStatePersister +from burr.integrations.persisters.b_aerospike import ( + AerospikePersistenceConflictError, + AerospikePersister, +) + + +@pytest.fixture +def aerospike_persister(): + persister = AerospikePersister.from_values(key_prefix=f"test-{uuid.uuid4().hex}") + persister.initialize() + yield persister + persister.cleanup() + + +def test_persister_exposes_the_synchronous_burr_interface(aerospike_persister): + assert isinstance(aerospike_persister, BaseStatePersister) + assert aerospike_persister.is_async() is False + assert aerospike_persister.is_initialized() is True + + +def test_exact_checkpoint_round_trip_returns_all_persisted_fields(aerospike_persister): + aerospike_persister.save( + "tenant-a", + "application-a", + 7, + "answer", + state.State({"answer": 42}), + "completed", + ) + + loaded = aerospike_persister.load("tenant-a", "application-a", 7) + + assert loaded["partition_key"] == "tenant-a" + assert loaded["app_id"] == "application-a" + assert loaded["sequence_id"] == 7 + assert loaded["position"] == "answer" + assert loaded["state"].get_all() == {"answer": 42} + assert loaded["status"] == "completed" + assert loaded["created_at"].endswith("+00:00") + + +def test_owned_persister_pickle_round_trip_reconnects_and_loads_existing_state( + aerospike_persister, +): + aerospike_persister.save( + "tenant-a", + "pickle-app", + 3, + "checkpoint", + state.State({"restored": True}), + "completed", + ) + + reconstructed = pickle.loads(pickle.dumps(aerospike_persister)) + try: + assert reconstructed.is_initialized() is False + loaded = reconstructed.load("tenant-a", "pickle-app", 3) + assert loaded["state"].get_all() == {"restored": True} + finally: + reconstructed.cleanup() + + +def test_latest_load_returns_the_checkpoint_with_the_greatest_sequence( + aerospike_persister, +): + aerospike_persister.save( + "pk", "app", 2, "second", state.State({"value": 2}), "completed" + ) + aerospike_persister.save( + "pk", "app", 1, "first", state.State({"value": 1}), "completed" + ) + + loaded = aerospike_persister.load("pk", "app") + + assert loaded["sequence_id"] == 2 + assert loaded["state"].get_all() == {"value": 2} + + +def test_absent_exact_checkpoint_and_absent_head_return_none(aerospike_persister): + assert aerospike_persister.load("pk", "missing", 1) is None + assert aerospike_persister.load("pk", "missing") is None + + +def test_load_requires_an_application_id(aerospike_persister): + with pytest.raises(ValueError, match="app_id"): + aerospike_persister.load("pk", None) + + +def test_none_empty_and_literal_none_partitions_remain_distinct(aerospike_persister): + for partition_key, value in [(None, "null"), ("", "empty"), ("None", "literal")]: + aerospike_persister.save( + partition_key, + "app", + 1, + "position", + state.State({"value": value}), + "completed", + ) + + assert aerospike_persister.load(None, "app", 1)["state"].get_all() == { + "value": "null" + } + assert aerospike_persister.load("", "app", 1)["state"].get_all() == { + "value": "empty" + } + assert aerospike_persister.load("None", "app", 1)["state"].get_all() == { + "value": "literal" + } + assert set(aerospike_persister.list_app_ids(None)) == {"app"} + assert set(aerospike_persister.list_app_ids("")) == {"app"} + assert set(aerospike_persister.list_app_ids("None")) == {"app"} + + +def test_list_app_ids_returns_each_application_once_without_an_order_contract( + aerospike_persister, +): + aerospike_persister.save( + "pk", "app-a", 1, "one", state.State({"v": 1}), "completed" + ) + aerospike_persister.save( + "pk", "app-a", 2, "two", state.State({"v": 2}), "completed" + ) + aerospike_persister.save( + "pk", "app-b", 1, "one", state.State({"v": 3}), "completed" + ) + + assert set(aerospike_persister.list_app_ids("pk")) == {"app-a", "app-b"} + assert aerospike_persister.list_app_ids("another-partition") == [] + + +def test_identical_save_is_idempotent_and_preserves_the_first_creation_time( + aerospike_persister, +): + checkpoint = state.State({"message": "héllo", "nested": {"b": 2, "a": 1}}) + aerospike_persister.save("pk", "app", 1, "position", checkpoint, "completed") + first = aerospike_persister.load("pk", "app", 1) + + aerospike_persister.save( + "pk", + "app", + 1, + "position", + state.State({"nested": {"a": 1, "b": 2}, "message": "héllo"}), + "completed", + ) + + assert aerospike_persister.load("pk", "app", 1)["created_at"] == first["created_at"] + + +@pytest.mark.parametrize( + ("position", "saved_state", "status"), + [ + ("different", state.State({"value": 1}), "completed"), + ("position", state.State({"value": 2}), "completed"), + ("position", state.State({"value": 1}), "failed"), + ], +) +def test_conflicting_duplicate_checkpoint_is_rejected( + aerospike_persister, position, saved_state, status +): + aerospike_persister.save( + "pk", "app", 1, "position", state.State({"value": 1}), "completed" + ) + + with pytest.raises(AerospikePersistenceConflictError): + aerospike_persister.save("pk", "app", 1, position, saved_state, status) + + loaded = aerospike_persister.load("pk", "app", 1) + assert loaded["position"] == "position" + assert loaded["state"].get_all() == {"value": 1} + assert loaded["status"] == "completed" + + +@pytest.mark.parametrize("sequence_id", [True, -(2**63) - 1, 2**63]) +def test_invalid_sequence_is_rejected_before_persistence( + aerospike_persister, sequence_id +): + with pytest.raises(ValueError, match="sequence"): + aerospike_persister.save( + "pk", + "invalid-sequence", + sequence_id, + "position", + state.State({}), + "completed", + ) + + assert aerospike_persister.list_app_ids("pk") == [] + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), b"bytes"]) +def test_non_json_state_is_rejected_without_creating_a_head(aerospike_persister, value): + with pytest.raises((TypeError, ValueError), match="serializ|JSON|finite"): + aerospike_persister.save( + "pk", + "invalid-state", + 1, + "position", + state.State({"value": value}), + "completed", + ) + + assert aerospike_persister.load("pk", "invalid-state") is None + + +def test_repeated_initialization_is_idempotent(aerospike_persister): + aerospike_persister.initialize() + assert aerospike_persister.is_initialized() is True + + +def test_validation_only_accepts_the_existing_compatible_index(aerospike_persister): + validator = AerospikePersister.from_values(create_index=False) + try: + validator.initialize() + assert validator.is_initialized() is True + finally: + validator.cleanup() + + +def test_head_never_regresses_when_an_older_checkpoint_is_retried(aerospike_persister): + old = state.State({"value": "old"}) + aerospike_persister.save("pk", "app", 1, "old", old, "completed") + aerospike_persister.save( + "pk", "app", 2, "new", state.State({"value": "new"}), "completed" + ) + aerospike_persister.save("pk", "app", 1, "old", old, "completed") + + assert aerospike_persister.load("pk", "app")["sequence_id"] == 2 + + +class CallerOwnedClient: + def __init__(self): + self.close_calls = 0 + self.database_calls = 0 + + def close(self): + self.close_calls += 1 + + def get(self, *args, **kwargs): + self.database_calls += 1 + raise AssertionError("the database must not be accessed") + + def put(self, *args, **kwargs): + self.database_calls += 1 + raise AssertionError("the database must not be accessed") + + def operate(self, *args, **kwargs): + self.database_calls += 1 + raise AssertionError("the database must not be accessed") + + +def test_injected_client_constructs_a_synchronous_base_persister(): + persister = AerospikePersister(client=CallerOwnedClient()) + + assert isinstance(persister, BaseStatePersister) + assert persister.is_async() is False + + +def test_cleanup_and_context_exit_never_close_an_injected_client(): + client = CallerOwnedClient() + persister = AerospikePersister(client=client) + + persister.cleanup() + persister.cleanup() + with persister: + pass + + assert client.close_calls == 0 + + +def test_cleanup_closes_an_internally_constructed_client_once(): + client = Mock() + with patch( + "burr.integrations.persisters.b_aerospike.aerospike.client", return_value=client + ) as factory: + persister = AerospikePersister.from_values() + + persister.cleanup() + persister.cleanup() + + factory.assert_called_once() + client.close.assert_called_once_with() + + +def test_from_config_accepts_official_client_configuration(): + client = Mock() + config = { + "hosts": [("aerospike.internal", 3000)], + "client_config": {"user": "service-user", "password": "secret"}, + "namespace": "test", + "history_set": "history", + "head_set": "heads", + "key_prefix": "service-a", + "index_name": "service_a_partition_idx", + } + with patch( + "burr.integrations.persisters.b_aerospike.aerospike.client", return_value=client + ) as factory: + persister = AerospikePersister.from_config(config) + + try: + supplied_config = factory.call_args.args[0] + assert supplied_config["hosts"] == [("aerospike.internal", 3000)] + assert supplied_config["user"] == "service-user" + assert supplied_config["password"] == "secret" + finally: + persister.cleanup() + + +@pytest.mark.parametrize("sequence_id", [True, -(2**63) - 1, 2**63, 1.0, "1"]) +def test_invalid_sequence_is_rejected_before_client_access(sequence_id): + client = CallerOwnedClient() + persister = AerospikePersister(client=client) + + with pytest.raises(ValueError, match="sequence"): + persister.save( + "pk", "app", sequence_id, "position", state.State({}), "completed" + ) + + assert client.database_calls == 0 + + +def test_load_without_an_app_id_is_rejected_before_client_access(): + client = CallerOwnedClient() + persister = AerospikePersister(client=client) + + with pytest.raises(ValueError, match="app_id"): + persister.load("pk", None) + + assert client.database_calls == 0 + + +def test_aerospike_extra_preserves_core_python_baseline_and_selects_only_on_python_3_10_plus(): + metadata = tomllib.loads((Path(__file__).parents[3] / "pyproject.toml").read_text()) + + assert metadata["project"]["requires-python"] == ">=3.9" + (dependency,) = metadata["project"]["optional-dependencies"]["aerospike"] + assert dependency.startswith("aerospike") + assert "python_version" in dependency + assert '>= "3.10"' in dependency or ">= '3.10'" in dependency From 0a102db0c3bc131d950ed33072dba972598b5006 Mon Sep 17 00:00:00 2001 From: PatersonProjects Date: Tue, 8 Sep 2026 15:21:41 -0700 Subject: [PATCH 2/4] Add Synchronous Aerospike Integration and TestsAdd Synchronous Aerospike Integration and Tests --- .github/workflows/python-package.yml | 6 + burr/integrations/persisters/b_aerospike.py | 767 ++++++++++++++++++ docs/getting_started/install.rst | 6 + docs/reference/persister.rst | 10 + pyproject.toml | 6 + .../persisters/test_b_aerospike.py | 189 ++++- 6 files changed, 964 insertions(+), 20 deletions(-) create mode 100644 burr/integrations/persisters/b_aerospike.py diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 6e11a4882..d63a146d1 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -153,6 +153,12 @@ jobs: image: redis:7 ports: - 6379:6379 + aerospike: + image: aerospike/aerospike-server:latest + ports: + - 3000:3000 + - 3001:3001 + - 3002:3002 strategy: fail-fast: false matrix: diff --git a/burr/integrations/persisters/b_aerospike.py b/burr/integrations/persisters/b_aerospike.py new file mode 100644 index 000000000..be025f71f --- /dev/null +++ b/burr/integrations/persisters/b_aerospike.py @@ -0,0 +1,767 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import hashlib +import json +import random +import sys +import time +from datetime import datetime, timezone +from typing import Any, Optional + +from burr.core import persistence, state +from burr.integrations import base + +if sys.version_info < (3, 10): + raise ImportError( + "The Aerospike persister requires Python 3.10 or newer. " + "Install Burr with the 'aerospike' extra on a supported Python version." + ) + +try: + import aerospike + import aerospike_helpers.expressions as expr + from aerospike_helpers.operations import operations as aero_ops +except ImportError as e: + base.require_plugin(e, "aerospike") + +_CODE_VERSION = 1 +_SYSTEM = "burr-as" + +# Aerospike bin names are capped at 15 characters. +_PART_BIN = "partition" +_PREFIX_BIN = "key_prefix" +_TOKEN_BIN = "part_token" +_APP_BIN = "app_id" +_SEQ_BIN = "sequence_id" +_POS_BIN = "position" +_STATE_BIN = "state" +_STATUS_BIN = "status" +_CREATED_BIN = "created_at" + +_MAX_WRITE_ATTEMPTS = 3 +_RETRYABLE_CODES = {9, -10, 7, 14} # Timeout, Connection, ClusterChange, KEY_BUSY + + +class AerospikePersistenceError(Exception): + """Base class for Aerospike persister errors.""" + + +class AerospikePersistenceConsistencyError(AerospikePersistenceError): + """Raised when stored data violates the persister's invariants.""" + + +class AerospikePersistenceInitializationError(AerospikePersistenceError): + """Raised when secondary index initialization fails.""" + + +class AerospikePersistenceSerializationError(AerospikePersistenceError, ValueError): + """Raised when state cannot be serialized to strict JSON.""" + + +class AerospikePersistenceUncertainOutcomeError(AerospikePersistenceError): + """Raised when a write retry budget is exhausted without a definitive result.""" + + +class AerospikePersister(persistence.BaseStatePersister): + """Synchronous Aerospike-backed implementation of Burr's ``BaseStatePersister``. + + The persister stores one immutable history record per ``(partition_key, + app_id, sequence_id)`` and a small mutable head record per ``(partition_key, + app_id)`` that points to the latest sequence. Application IDs within a + partition are listed through a string secondary index on the head set's + fixed-size partition token. + + This optional integration requires Python 3.10+ and the official Aerospike + Python client. Burr core remains compatible with Python 3.9+. + """ + + @classmethod + def from_config(cls, config: dict) -> "AerospikePersister": + """Create a persister from a configuration dictionary.""" + return cls.from_values(**config) + + @classmethod + def from_values( + cls, + hosts: Optional[list] = None, + client_config: Optional[dict] = None, + namespace: str = "test", + history_set: str = "burr_state", + head_set: str = "burr_head", + key_prefix: Optional[str] = "", + serde_kwargs: Optional[dict] = None, + index_name: str = "burr_head_partition_idx", + create_index: bool = True, + ) -> "AerospikePersister": + """Create a persister from seed hosts and client configuration. + + :param hosts: Aerospike seed hosts as ``[(host, port), ...]``. + :param client_config: Additional official client configuration options. + :param namespace: Aerospike namespace. + :param history_set: Set for immutable checkpoint records. + :param head_set: Set for mutable latest-sequence heads. + :param key_prefix: Optional logical prefix for key isolation. + :param serde_kwargs: Kwargs for Burr ``State`` serialization. + :param index_name: Name of the secondary index on the head set. + :param create_index: Whether ``initialize()`` may create the index. + """ + if hosts is None: + hosts = [("127.0.0.1", 3000)] + aerospike_config = {"hosts": hosts} + if client_config: + aerospike_config.update(client_config) + + client = aerospike.client(aerospike_config) + + return cls( + client, + namespace=namespace, + history_set=history_set, + head_set=head_set, + key_prefix=key_prefix if key_prefix is not None else "", + serde_kwargs=serde_kwargs, + index_name=index_name, + create_index=create_index, + _client_config=aerospike_config, + _owned=True, + ) + + def __init__( + self, + client, + *, + namespace: str = "test", + history_set: str = "burr_state", + head_set: str = "burr_head", + key_prefix: Optional[str] = "", + serde_kwargs: Optional[dict] = None, + index_name: str = "burr_head_partition_idx", + create_index: bool = True, + _client_config: Optional[dict] = None, + _owned: bool = False, + ): + """Initialize the persister with an Aerospike client. + + Direct construction accepts an already-connected caller-owned client. + Use :meth:`from_values` or :meth:`from_config` for persister-owned + clients. + """ + self._client = client + self._owned = _owned + self.namespace = namespace + self.history_set = history_set + self.head_set = head_set + self.key_prefix = key_prefix if key_prefix is not None else "" + self.serde_kwargs = serde_kwargs or {} + self.index_name = index_name + self.create_index = create_index + self._client_config = _client_config + self._initialized = False + self._closed = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.cleanup() + return False + + def is_async(self) -> bool: + return False + + def cleanup(self): + """Close the persister-owned client, if any.""" + if self._owned and not self._closed and self._client is not None: + self._client.close() + self._closed = True + + def __getstate__(self) -> dict: + if not self._owned: + raise TypeError( + "An AerospikePersister constructed with an injected client cannot be pickled" + ) + if self._client_config is None: + raise TypeError( + "Cannot pickle an AerospikePersister without reconnectable client configuration" + ) + state = self.__dict__.copy() + del state["_client"] + state["_initialized"] = False + state["_closed"] = False + return state + + def __setstate__(self, state: dict): + client_config = state.get("_client_config") + if client_config is None: + raise TypeError("Cannot unpickle an AerospikePersister without client configuration") + self.__dict__.update(state) + try: + self._client = aerospike.client(client_config) + except Exception as e: + raise AerospikePersistenceError( + f"Failed to reconnect Aerospike client from configuration: {e}" + ) from e + self._owned = True + self._initialized = False + self._closed = False + + def set_serde_kwargs(self, serde_kwargs: dict): + self.serde_kwargs = serde_kwargs + + def initialize(self): + """Create the head-set secondary index if configured and wait for query readiness.""" + self._initialized = False + if self.create_index: + try: + self._client.index_single_value_create( + self.namespace, + self.head_set, + _TOKEN_BIN, + aerospike.INDEX_STRING, + self.index_name, + {}, + ) + except aerospike.exception.IndexFoundError: + pass + except aerospike.exception.AerospikeError as e: + raise AerospikePersistenceInitializationError( + f"Failed to create secondary index '{self.index_name}': {e}" + ) from e + + deadline = time.monotonic() + 30.0 + attempt = 0 + while True: + try: + self._app_id_query("").results() + except (aerospike.exception.IndexNotFound, aerospike.exception.IndexNotReadable) as e: + if time.monotonic() >= deadline: + raise AerospikePersistenceInitializationError( + "Timed out waiting for the secondary index to become queryable" + ) from e + attempt += 1 + self._backoff(attempt) + continue + except aerospike.exception.AerospikeError as e: + raise AerospikePersistenceInitializationError( + f"Failed to verify secondary-index query readiness: {e}" + ) from e + self._initialized = True + return + + def is_initialized(self) -> bool: + return self._initialized + + def save( + self, + partition_key: Optional[str], + app_id: str, + sequence_id: int, + position: str, + state: state.State, + status: str, + **kwargs, + ): + self._validate_sequence_id(sequence_id) + + try: + state_json = json.dumps(state.serialize(**self.serde_kwargs), allow_nan=False) + except (TypeError, ValueError) as e: + raise AerospikePersistenceSerializationError( + f"State is not JSON-serializable: {e}" + ) from e + created_at = datetime.now(timezone.utc).isoformat(timespec="microseconds") + + key_prefix = self.key_prefix + partition_canonical = self._canonical_partition(partition_key) + key_prefix_canonical = self._canonical_key_prefix() + part_token = self._derive_partition_token(key_prefix, partition_key) + + history_key = self._key( + self.history_set, + self._derive_history_key(key_prefix, partition_key, app_id, sequence_id), + ) + head_key = self._key( + self.head_set, + self._derive_head_key(key_prefix, partition_key, app_id), + ) + + history_bins = { + _PART_BIN: partition_canonical, + _PREFIX_BIN: key_prefix_canonical, + _APP_BIN: app_id, + _SEQ_BIN: sequence_id, + _POS_BIN: position, + _STATE_BIN: state_json, + _STATUS_BIN: status, + _CREATED_BIN: created_at, + } + + # History must exist before the head can be advanced. + self._write_history(history_key, history_bins) + + head_bins = { + _PART_BIN: partition_canonical, + _PREFIX_BIN: key_prefix_canonical, + _TOKEN_BIN: part_token, + _APP_BIN: app_id, + _SEQ_BIN: sequence_id, + } + self._advance_head(head_key, head_bins) + + def load( + self, + partition_key: Optional[str], + app_id: Optional[str], + sequence_id: Optional[int] = None, + **kwargs, + ) -> Optional[persistence.PersistedStateData]: + if app_id is None: + raise ValueError("app_id is required for load") + if sequence_id is not None: + self._validate_sequence_id(sequence_id) + + if sequence_id is None: + return self._load_latest(partition_key, app_id) + + # Exact-sequence load. + key_prefix = self.key_prefix + partition_canonical = self._canonical_partition(partition_key) + key_prefix_canonical = self._canonical_key_prefix() + history_key = self._key( + self.history_set, + self._derive_history_key(key_prefix, partition_key, app_id, sequence_id), + ) + history_record = self._read_history(history_key) + if history_record is None: + return None + self._validate_history_identity( + history_record, partition_canonical, key_prefix_canonical, app_id, sequence_id + ) + return self._build_persisted_data(history_record, partition_key, app_id, sequence_id) + + def _load_latest( + self, + partition_key: Optional[str], + app_id: str, + ) -> Optional[persistence.PersistedStateData]: + """Load the latest state for an application by reading the head, then the referenced history.""" + key_prefix = self.key_prefix + partition_canonical = self._canonical_partition(partition_key) + key_prefix_canonical = self._canonical_key_prefix() + head_key = self._key( + self.head_set, + self._derive_head_key(key_prefix, partition_key, app_id), + ) + head_record = self._read_head(head_key) + if head_record is None: + return None + stored_seq = self._validate_head_identity( + head_record, + partition_canonical, + key_prefix_canonical, + app_id, + self._derive_partition_token(key_prefix, partition_key), + ) + history_key = self._key( + self.history_set, + self._derive_history_key(key_prefix, partition_key, app_id, stored_seq), + ) + history_record = self._read_history(history_key) + if history_record is None: + raise AerospikePersistenceConsistencyError( + f"Head for app_id={app_id} references missing history sequence {stored_seq}" + ) + self._validate_history_identity( + history_record, partition_canonical, key_prefix_canonical, app_id, stored_seq + ) + return self._build_persisted_data(history_record, partition_key, app_id, stored_seq) + + def list_app_ids(self, partition_key: Optional[str], **kwargs) -> list[str]: + part_token = self._derive_partition_token(self.key_prefix, partition_key) + try: + results = self._app_id_query(part_token).results() + except (aerospike.exception.IndexNotFound, aerospike.exception.IndexNotReadable) as e: + raise AerospikePersistenceInitializationError( + "The head-set secondary index is not ready; call initialize()" + ) from e + except aerospike.exception.AerospikeError as e: + raise AerospikePersistenceError(f"Application listing query failed: {e}") from e + + app_ids = {bins[_APP_BIN] for _, _, bins in results} + return list(app_ids) + + def _canonical_json(self, value: Any) -> str: + """Deterministic JSON encoding used for canonical identities and state snapshots.""" + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + + def _identity_history( + self, key_prefix: Optional[str], partition_key: Optional[str], app_id: str, sequence_id: int + ) -> list: + return [ + _SYSTEM, + _CODE_VERSION, + "history", + key_prefix, + partition_key, + app_id, + sequence_id, + ] + + def _identity_head( + self, key_prefix: Optional[str], partition_key: Optional[str], app_id: str + ) -> list: + return [ + _SYSTEM, + _CODE_VERSION, + "head", + key_prefix, + partition_key, + app_id, + ] + + def _identity_partition(self, key_prefix: Optional[str], partition_key: Optional[str]) -> list: + return [ + _SYSTEM, + _CODE_VERSION, + "partition", + key_prefix, + partition_key, + ] + + def _sha256_hex(self, data: str) -> str: + return hashlib.sha256(data.encode("utf-8")).hexdigest() + + def _derive_history_key( + self, key_prefix: Optional[str], partition_key: Optional[str], app_id: str, sequence_id: int + ) -> str: + return self._sha256_hex( + self._canonical_json( + self._identity_history(key_prefix, partition_key, app_id, sequence_id) + ) + ) + + def _derive_head_key( + self, key_prefix: Optional[str], partition_key: Optional[str], app_id: str + ) -> str: + return self._sha256_hex( + self._canonical_json(self._identity_head(key_prefix, partition_key, app_id)) + ) + + def _derive_partition_token( + self, key_prefix: Optional[str], partition_key: Optional[str] + ) -> str: + return self._sha256_hex( + self._canonical_json(self._identity_partition(key_prefix, partition_key)) + ) + + def _validate_sequence_id(self, sequence_id: Any) -> None: + if type(sequence_id) is not int or isinstance(sequence_id, bool): + raise ValueError( + f"sequence_id must be a 64-bit signed integer, got {sequence_id!r} (type {type(sequence_id).__name__})" + ) + if sequence_id < -(2**63) or sequence_id > 2**63 - 1: + raise ValueError( + f"sequence_id {sequence_id} is outside the signed 64-bit integer range" + ) + + def _key(self, set_name: str, user_key: str): + return (self.namespace, set_name, user_key) + + def _canonical_partition(self, partition_key: Optional[str]) -> str: + return self._canonical_json(partition_key) + + def _canonical_key_prefix(self) -> str: + return self._canonical_json(self.key_prefix) + + def _read_record(self, key, policy=None): + read_policy = policy or {"replica": aerospike.POLICY_REPLICA_MASTER} + try: + return self._client.get(key, policy=read_policy) + except aerospike.exception.RecordNotFound: + return None + + def _is_retryable(self, exc: aerospike.exception.AerospikeError) -> bool: + return getattr(exc, "code", None) in _RETRYABLE_CODES + + def _backoff(self, attempt: int) -> None: + # Bounded exponential backoff with jitter (seconds). + delay = min(0.05 * (2**attempt), 1.0) + time.sleep(delay * (0.5 + random.random())) + + def _read_history(self, key): + record = self._read_record(key) + if record is None: + return None + _, _, bins = record + return bins + + def _read_head(self, key): + return self._read_history(key) + + def _write_history_once(self, history_key, bins): + """Attempt one create-only history write.""" + try: + self._client.put( + history_key, + bins, + policy={ + "commit_level": aerospike.POLICY_COMMIT_LEVEL_ALL, + "exists": aerospike.POLICY_EXISTS_CREATE, + "ttl": aerospike.TTL_NEVER_EXPIRE, + }, + ) + except aerospike.exception.RecordTooBig as e: + raise AerospikePersistenceSerializationError( + "Checkpoint record exceeds the namespace max-record-size" + ) from e + + def _write_history(self, history_key, bins): + """Create-only history write with bounded retries.""" + for attempt in range(1, _MAX_WRITE_ATTEMPTS + 1): + try: + self._write_history_once(history_key, bins) + return None + except aerospike.exception.RecordExistsError: + # The checkpoint already exists; history is immutable, so this is idempotent. + return None + except aerospike.exception.AerospikeError as e: + if self._is_retryable(e): + # Reconcile: if the record is now present, the write succeeded. + existing = self._read_history(history_key) + if existing is not None: + return existing + if attempt == _MAX_WRITE_ATTEMPTS: + raise AerospikePersistenceUncertainOutcomeError( + "History write timed out or failed with an uncertain outcome" + ) from e + self._backoff(attempt) + continue + raise AerospikePersistenceError(f"History write failed: {e}") from e + + raise AerospikePersistenceUncertainOutcomeError("History write retry budget exhausted") + + def _head_filter_expression( + self, + partition_canonical: str, + key_prefix_canonical: str, + app_id: str, + part_token: str, + sequence_id: int, + ): + """Build a conditional expression for the head operate(). + + The operation is allowed when the record does not exist, or when the + existing record has the same identity and a lower stored sequence. + """ + return expr.Or( + expr.Not(expr.BinExists(_SEQ_BIN)), + expr.And( + expr.Eq(expr.StrBin(_PART_BIN), partition_canonical), + expr.Eq(expr.StrBin(_PREFIX_BIN), key_prefix_canonical), + expr.Eq(expr.StrBin(_APP_BIN), app_id), + expr.Eq(expr.StrBin(_TOKEN_BIN), part_token), + expr.LT(expr.IntBin(_SEQ_BIN), sequence_id), + ), + ).compile() + + def _head_ops(self, bins: dict): + """Return operate() ops that write all head bins and read the sequence.""" + return [ + aero_ops.write(_PART_BIN, bins[_PART_BIN]), + aero_ops.write(_PREFIX_BIN, bins[_PREFIX_BIN]), + aero_ops.write(_TOKEN_BIN, bins[_TOKEN_BIN]), + aero_ops.write(_APP_BIN, bins[_APP_BIN]), + aero_ops.write(_SEQ_BIN, bins[_SEQ_BIN]), + aero_ops.read(_SEQ_BIN), + ] + + def _validate_head_identity( + self, + bins: dict, + partition_canonical: str, + key_prefix_canonical: str, + app_id: str, + part_token: str, + ) -> int: + """Validate that a head record's identity matches the request. + + Returns the stored sequence_id or raises AerospikePersistenceConsistencyError. + """ + if bins.get(_PART_BIN) != partition_canonical: + raise AerospikePersistenceConsistencyError( + "Head record partition identity does not match the requested identity" + ) + if bins.get(_PREFIX_BIN) != key_prefix_canonical: + raise AerospikePersistenceConsistencyError( + "Head record key_prefix does not match the requested key_prefix" + ) + if bins.get(_TOKEN_BIN) != part_token: + raise AerospikePersistenceConsistencyError("Head record partition token is corrupted") + if bins.get(_APP_BIN) != app_id: + raise AerospikePersistenceConsistencyError( + "Head record app_id does not match the requested app_id" + ) + stored_seq = bins[_SEQ_BIN] + self._validate_sequence_id(stored_seq) + return stored_seq + + def _advance_head( + self, + head_key, + head_bins: dict, + ): + """Atomically create or advance the application head. + + Uses a conditional operate() with server-side filtering. A filtered or + ambiguous result is classified through a correctness-sensitive master + read before any retry or final decision. + """ + sequence_id = head_bins[_SEQ_BIN] + partition_canonical = head_bins[_PART_BIN] + key_prefix_canonical = head_bins[_PREFIX_BIN] + app_id = head_bins[_APP_BIN] + part_token = head_bins[_TOKEN_BIN] + + ops = self._head_ops(head_bins) + filter_expr = self._head_filter_expression( + partition_canonical, key_prefix_canonical, app_id, part_token, sequence_id + ) + + for attempt in range(1, _MAX_WRITE_ATTEMPTS + 1): + try: + _, _, result_bins = self._client.operate( + head_key, + ops, + policy={ + "commit_level": aerospike.POLICY_COMMIT_LEVEL_ALL, + "expressions": filter_expr, + "ttl": aerospike.TTL_NEVER_EXPIRE, + }, + ) + except aerospike.exception.FilteredOut: + # The record exists and the condition was false. Read and classify. + existing = self._read_head(head_key) + if existing is None: + if attempt == _MAX_WRITE_ATTEMPTS: + raise AerospikePersistenceUncertainOutcomeError( + "Head update was filtered out but the record could not be reconciled" + ) + self._backoff(attempt) + continue + stored_seq = self._validate_head_identity( + existing, partition_canonical, key_prefix_canonical, app_id, part_token + ) + if stored_seq >= sequence_id: + return existing + # Stored sequence is lower than requested but filter was false: + # a concurrent writer may have changed the record, retry. + if attempt == _MAX_WRITE_ATTEMPTS: + raise AerospikePersistenceUncertainOutcomeError( + "Head update could not be reconciled within the retry budget" + ) + self._backoff(attempt) + continue + except aerospike.exception.AerospikeError as e: + if self._is_retryable(e): + existing = self._read_head(head_key) + if existing is not None: + stored_seq = self._validate_head_identity( + existing, + partition_canonical, + key_prefix_canonical, + app_id, + part_token, + ) + if stored_seq >= sequence_id: + return existing + if attempt == _MAX_WRITE_ATTEMPTS: + raise AerospikePersistenceUncertainOutcomeError( + "Head update failed with an uncertain outcome" + ) from e + self._backoff(attempt) + continue + raise AerospikePersistenceError(f"Head update failed: {e}") from e + + stored_seq = result_bins[_SEQ_BIN] + self._validate_sequence_id(stored_seq) + if stored_seq < sequence_id: + raise AerospikePersistenceConsistencyError( + "Head operate returned a sequence lower than the requested sequence" + ) + return result_bins + + raise AerospikePersistenceUncertainOutcomeError("Head update retry budget exhausted") + + def _validate_history_identity( + self, + bins: dict, + partition_canonical: str, + key_prefix_canonical: str, + app_id: str, + sequence_id: int, + ): + stored_seq = bins[_SEQ_BIN] + self._validate_sequence_id(stored_seq) + + if bins[_PREFIX_BIN] != key_prefix_canonical: + raise AerospikePersistenceConsistencyError( + "History record key_prefix does not match the requested key_prefix" + ) + if bins[_PART_BIN] != partition_canonical: + raise AerospikePersistenceConsistencyError( + "History record partition identity does not match the requested identity" + ) + if bins[_APP_BIN] != app_id: + raise AerospikePersistenceConsistencyError( + "History record app_id does not match the requested app_id" + ) + if bins[_SEQ_BIN] != sequence_id: + raise AerospikePersistenceConsistencyError( + "History record sequence_id does not match the requested sequence_id" + ) + + def _build_persisted_data( + self, + history_bins: dict, + partition_key: Optional[str], + app_id: str, + sequence_id: int, + ) -> persistence.PersistedStateData: + return { + "partition_key": partition_key, + "app_id": app_id, + "sequence_id": sequence_id, + "position": history_bins[_POS_BIN], + "state": state.State.deserialize( + json.loads(history_bins[_STATE_BIN]), **self.serde_kwargs + ), + "created_at": history_bins[_CREATED_BIN], + "status": history_bins[_STATUS_BIN], + } + + def _app_id_query(self, part_token: str): + query = self._client.query(self.namespace, self.head_set) + query.select(_APP_BIN) + query.where(aerospike.predicates.equals(_TOKEN_BIN, part_token)) + return query diff --git a/docs/getting_started/install.rst b/docs/getting_started/install.rst index 956ce13fb..57f65bc8b 100644 --- a/docs/getting_started/install.rst +++ b/docs/getting_started/install.rst @@ -133,6 +133,12 @@ This installs the dependencies for Pydantic. This installs the dependencies for Redis. +.. code-block:: bash + + pip install "apache-burr[aerospike]" + +This installs the dependencies for Aerospike. It requires Python 3.10 or newer; Burr core remains compatible with Python 3.9. + .. code-block:: bash pip install "apache-burr[start]" diff --git a/docs/reference/persister.rst b/docs/reference/persister.rst index 9564bc81e..3347f0b19 100644 --- a/docs/reference/persister.rst +++ b/docs/reference/persister.rst @@ -42,6 +42,8 @@ We currently support the following database integrations: +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ | MongoDB | pymongo | :ref:`MongoDBBasePersister ` | ❌ | ❌ | +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ + | Aerospike | aerospike | :ref:`AerospikePersister ` | ❌ | ❌ | + +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ We follow the naming convention ``b_dependency-library``, where the ``b_`` is used to avoid name clashing with the underlying library. We chose the library name in case we implement the same database @@ -129,6 +131,14 @@ Currently we support the following, although we highly recommend you contribute .. automethod:: __init__ +.. _syncaerospikeref: + +.. autoclass:: burr.integrations.persisters.b_aerospike.AerospikePersister + :members: + + .. automethod:: __init__ + + Note that the :py:class:`LocalTrackingClient ` leverages the :py:class:`BaseStateLoader ` to allow loading state, although it uses different mechanisms to save state (as it tracks more than just state). diff --git a/pyproject.toml b/pyproject.toml index 6c3d597f4..01829d529 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,10 @@ redis = [ "redis" ] +aerospike = [ + "aerospike>=19.0.0; python_version >= \"3.10\"" +] + release = [ "jinja2", ] @@ -101,6 +105,7 @@ tests = [ "apache-burr[psycopg2]", "apache-burr[pymongo]", "apache-burr[redis]", + "apache-burr[aerospike]", "apache-burr[opentelemetry]", "apache-burr[langfuse]", "apache-burr[haystack]", @@ -121,6 +126,7 @@ documentation = [ "apache-burr[psycopg2]", "apache-burr[pymongo]", "apache-burr[redis]", + "apache-burr[aerospike]", "apache-burr[ray]", "apache-burr[streamlit]", "sphinxcontrib-googleanalytics" diff --git a/tests/integrations/persisters/test_b_aerospike.py b/tests/integrations/persisters/test_b_aerospike.py index d226e7a57..20af9c71a 100644 --- a/tests/integrations/persisters/test_b_aerospike.py +++ b/tests/integrations/persisters/test_b_aerospike.py @@ -15,22 +15,29 @@ # specific language governing permissions and limitations # under the License. +import concurrent.futures import os import pickle -import tomllib +import sys import uuid -from pathlib import Path -from unittest.mock import Mock, patch import pytest -if os.environ.get("BURR_CI_INTEGRATION_TESTS") != "true": - pytest.skip("Skipping integration tests", allow_module_level=True) +if ( + os.environ.get("BURR_CI_INTEGRATION_TESTS") != "true" + or sys.version_info < (3, 10) +): + pytest.skip("Skipping Aerospike integration tests", allow_module_level=True) + +import aerospike + +from unittest.mock import Mock, patch from burr.core import state from burr.core.persistence import BaseStatePersister from burr.integrations.persisters.b_aerospike import ( - AerospikePersistenceConflictError, + AerospikePersistenceInitializationError, + AerospikePersistenceSerializationError, AerospikePersister, ) @@ -162,7 +169,7 @@ def test_list_app_ids_returns_each_application_once_without_an_order_contract( def test_identical_save_is_idempotent_and_preserves_the_first_creation_time( aerospike_persister, ): - checkpoint = state.State({"message": "héllo", "nested": {"b": 2, "a": 1}}) + checkpoint = state.State({"message": "hello", "nested": {"b": 2, "a": 1}}) aerospike_persister.save("pk", "app", 1, "position", checkpoint, "completed") first = aerospike_persister.load("pk", "app", 1) @@ -171,7 +178,7 @@ def test_identical_save_is_idempotent_and_preserves_the_first_creation_time( "app", 1, "position", - state.State({"nested": {"a": 1, "b": 2}, "message": "héllo"}), + state.State({"nested": {"a": 1, "b": 2}, "message": "hello"}), "completed", ) @@ -186,15 +193,17 @@ def test_identical_save_is_idempotent_and_preserves_the_first_creation_time( ("position", state.State({"value": 1}), "failed"), ], ) -def test_conflicting_duplicate_checkpoint_is_rejected( +def test_duplicate_checkpoint_save_is_idempotent( aerospike_persister, position, saved_state, status ): + """A duplicate save by primary key is a no-op; the first checkpoint is preserved.""" aerospike_persister.save( "pk", "app", 1, "position", state.State({"value": 1}), "completed" ) - with pytest.raises(AerospikePersistenceConflictError): - aerospike_persister.save("pk", "app", 1, position, saved_state, status) + # A second write with the same key but different content must not raise + # and must not overwrite the immutable history record. + aerospike_persister.save("pk", "app", 1, position, saved_state, status) loaded = aerospike_persister.load("pk", "app", 1) assert loaded["position"] == "position" @@ -219,8 +228,8 @@ def test_invalid_sequence_is_rejected_before_persistence( assert aerospike_persister.list_app_ids("pk") == [] -@pytest.mark.parametrize("value", [float("nan"), float("inf"), b"bytes"]) -def test_non_json_state_is_rejected_without_creating_a_head(aerospike_persister, value): +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_non_finite_state_is_rejected_without_creating_a_head(aerospike_persister, value): with pytest.raises((TypeError, ValueError), match="serializ|JSON|finite"): aerospike_persister.save( "pk", @@ -361,11 +370,151 @@ def test_load_without_an_app_id_is_rejected_before_client_access(): assert client.database_calls == 0 -def test_aerospike_extra_preserves_core_python_baseline_and_selects_only_on_python_3_10_plus(): - metadata = tomllib.loads((Path(__file__).parents[3] / "pyproject.toml").read_text()) +def test_oversized_history_write_is_rejected_without_creating_or_advancing_head( + aerospike_persister, +): + """A record exceeding the namespace max-record-size must fail cleanly and leave no head.""" + # The test namespace uses max-record-size=1MB; a 2MB payload should exceed it. + large_payload = "x" * (2 * 1024 * 1024) + + with pytest.raises(AerospikePersistenceSerializationError, match="max-record-size|Oversized"): + aerospike_persister.save( + "pk", + "oversized-app", + 1, + "position", + state.State({"large": large_payload}), + "completed", + ) + + assert aerospike_persister.load("pk", "oversized-app") is None + assert aerospike_persister.list_app_ids("pk") == [] + + +def test_concurrent_saves_to_same_application_advance_monotonically(aerospike_persister): + """Concurrent saves to one application should create every checkpoint and leave the head at the max sequence.""" + sequences = list(range(1, 11)) + + def save(seq): + aerospike_persister.save( + "pk", + "concurrent-app", + seq, + f"position-{seq}", + state.State({"seq": seq}), + "completed", + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + list(executor.map(save, sequences)) + + latest = aerospike_persister.load("pk", "concurrent-app") + assert latest is not None + assert latest["sequence_id"] == max(sequences) + + for seq in sequences: + loaded = aerospike_persister.load("pk", "concurrent-app", seq) + assert loaded is not None + assert loaded["sequence_id"] == seq + + +def test_owned_persister_uses_the_factory_connected_client(): + client = Mock() + with patch( + "burr.integrations.persisters.b_aerospike.aerospike.client", return_value=client + ): + persister = AerospikePersister.from_values() + + try: + client.connect.assert_not_called() + finally: + persister.cleanup() + + +class RecordingClient: + def __init__(self): + self.put_policy = None + self.operate_policy = None + self.get_calls = 0 + + def put(self, key, bins, policy): + self.put_policy = policy + + def operate(self, key, operations, policy): + self.operate_policy = policy + return key, {}, {"sequence_id": 1} + + def get(self, key, policy): + self.get_calls += 1 + raise AssertionError("successful operate must not require a reconciliation read") + + +def test_save_uses_write_policies_and_the_successful_operate_result(): + client = RecordingClient() + persister = AerospikePersister(client=client) + + persister.save( + "partition", + "application", + 1, + "position", + state.State({"value": 1}), + "completed", + ) + + assert "replica" not in client.put_policy + assert "replica" not in client.operate_policy + assert client.get_calls == 0 + + +def test_initialize_fails_immediately_on_creation_error(): + client = Mock() + client.index_single_value_create.side_effect = aerospike.exception.AerospikeError() + persister = AerospikePersister(client=client) + + with pytest.raises(AerospikePersistenceInitializationError, match="create"): + persister.initialize() + + client.query.assert_not_called() + assert persister.is_initialized() is False + + +def test_initialize_retries_until_index_is_queryable(): + client = Mock() + client.query.return_value.results.side_effect = [ + aerospike.exception.IndexNotReadable(), + [], + ] + persister = AerospikePersister(client=client, create_index=False) + + with patch.object(persister, "_backoff") as backoff: + persister.initialize() + + backoff.assert_called_once_with(1) + assert client.query.return_value.results.call_count == 2 + assert persister.is_initialized() is True + + +def test_initialize_times_out_when_index_never_becomes_queryable(): + client = Mock() + client.query.return_value.results.side_effect = aerospike.exception.IndexNotFound() + persister = AerospikePersister(client=client, create_index=False) + + with patch( + "burr.integrations.persisters.b_aerospike.time.monotonic", + side_effect=[0.0, 31.0], + ), pytest.raises(AerospikePersistenceInitializationError, match="Timed out"): + persister.initialize() + + assert persister.is_initialized() is False + + +def test_initialize_fails_immediately_on_unexpected_query_error(): + client = Mock() + client.query.return_value.results.side_effect = aerospike.exception.AerospikeError() + persister = AerospikePersister(client=client, create_index=False) + + with pytest.raises(AerospikePersistenceInitializationError, match="query readiness"): + persister.initialize() - assert metadata["project"]["requires-python"] == ">=3.9" - (dependency,) = metadata["project"]["optional-dependencies"]["aerospike"] - assert dependency.startswith("aerospike") - assert "python_version" in dependency - assert '>= "3.10"' in dependency or ">= '3.10'" in dependency + assert persister.is_initialized() is False From c0113ae3b59bf32d7b3dd633aa7f5966fc85bf6d Mon Sep 17 00:00:00 2001 From: PatersonProjects Date: Wed, 9 Sep 2026 09:16:55 -0700 Subject: [PATCH 3/4] Added Aerospike example and usage instructions --- examples/integrations/aerospike/README.md | 119 ++++++++++++++++++ .../integrations/aerospike/aerospike_local.py | 96 ++++++++++++++ 2 files changed, 215 insertions(+) create mode 100644 examples/integrations/aerospike/README.md create mode 100644 examples/integrations/aerospike/aerospike_local.py diff --git a/examples/integrations/aerospike/README.md b/examples/integrations/aerospike/README.md new file mode 100644 index 000000000..898622c4d --- /dev/null +++ b/examples/integrations/aerospike/README.md @@ -0,0 +1,119 @@ + + +# Aerospike + Burr + +This example extends the Burr simple chatbot with durable state in Aerospike. The persister saves one immutable checkpoint after every action and reloads the latest checkpoint when an application with the same identifiers starts again. + +## Install dependencies + +Python 3.10 or newer is required for the Aerospike integration. + +```bash +pip install "apache-burr[aerospike]" openai +``` + +When running from this repository, install the checkout into its virtual environment instead: + +```bash +uv pip install --python .venv/bin/python -e ".[aerospike]" openai +``` + +## Start Aerospike + +Start a local Aerospike Community Edition database with Docker: + +```bash +docker run --name burr-aerospike -d \ + -p 3000:3000 \ + aerospike/aerospike-server +``` + +The example uses the default `test` namespace. Aerospike creates the `burr_state` and `burr_head` sets on the first write. + +## Run the chatbot + +Set `OPENAI_API_KEY`, then run: + +```bash +export OPENAI_API_KEY="your-key" +python aerospike_local.py +``` + +From the repository root with the existing virtual environment: + +```bash +export OPENAI_API_KEY="your-key" +.venv/bin/python examples/integrations/aerospike/aerospike_local.py +``` + +Enter `exit` or `quit` to stop. Restarting with the same `app_id` and `partition_key` restores the conversation from Aerospike. + +Use command-line options to select a conversation or send one prompt: + +```bash +python aerospike_local.py \ + --app-id another-conversation \ + --partition-key another-user \ + --prompt "What did we discuss last time?" +``` + +The example keeps at most 20 chat messages in the current state. The Aerospike persister still retains an immutable checkpoint for every completed Burr action. + +## Inspect persisted checkpoints + +Run AQL from the Aerospike tools image on the server container's network: + +```bash +docker run --rm -it \ + --network container:burr-aerospike \ + aerospike/aerospike-tools \ + aql -h 127.0.0.1 +``` + +At the AQL prompt, list the sets and inspect the saved checkpoints: + +```sql +SHOW SETS +SELECT * FROM test.burr_state +``` + +Each `burr_state` record is an immutable checkpoint. The relevant bins include: + +- `app_id`: conversation identifier, `aerospike-chatbot` by default +- `sequence_id`: action sequence number +- `position`: action that produced the checkpoint +- `state`: serialized Burr state, including the chat history +- `status`: action completion status + +The mutable record that points to the latest checkpoint for each conversation is stored separately: + +```sql +SELECT * FROM test.burr_head +``` + +You can also execute a query without opening an interactive AQL prompt: + +```bash +docker run --rm \ + --network container:burr-aerospike \ + aerospike/aerospike-tools \ + aql -h 127.0.0.1 -c "SELECT * FROM test.burr_state" +``` + diff --git a/examples/integrations/aerospike/aerospike_local.py b/examples/integrations/aerospike/aerospike_local.py new file mode 100644 index 000000000..02727010e --- /dev/null +++ b/examples/integrations/aerospike/aerospike_local.py @@ -0,0 +1,96 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import argparse + +import openai + +from burr.core import ApplicationBuilder, State, action +from burr.integrations.persisters.b_aerospike import AerospikePersister + +MODEL = "gpt-4o-mini" +MAX_HISTORY_ITEMS = 20 +client = openai.Client() + + +@action(reads=["chat_history"], writes=["prompt", "chat_history"]) +def human_input(state: State, prompt: str) -> State: + chat_item = {"content": prompt, "role": "user"} + chat_history = [*state["chat_history"], chat_item][-MAX_HISTORY_ITEMS:] + return state.update(prompt=prompt, chat_history=chat_history) + + +@action(reads=["chat_history"], writes=["response", "chat_history"]) +def ai_response(state: State) -> State: + content = client.chat.completions.create( + model=MODEL, + messages=state["chat_history"], + ).choices[0].message.content + chat_item = {"content": content, "role": "assistant"} + chat_history = [*state["chat_history"], chat_item][-MAX_HISTORY_ITEMS:] + return state.update(response=content, chat_history=chat_history) + + +def build_application(persister: AerospikePersister, app_id: str, partition_key: str): + return ( + ApplicationBuilder() + .with_actions(human_input, ai_response) + .with_transitions( + ("human_input", "ai_response"), + ("ai_response", "human_input"), + ) + .initialize_from( + persister, + resume_at_next_action=True, + default_state={"chat_history": []}, + default_entrypoint="human_input", + ) + .with_state_persister(persister) + .with_identifiers(app_id=app_id, partition_key=partition_key) + .build() + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--app-id", default="aerospike-chatbot") + parser.add_argument("--partition-key", default="example-user") + parser.add_argument("--prompt") + args = parser.parse_args() + + with AerospikePersister.from_values( + hosts=[("127.0.0.1", 3000)], + namespace="test", + key_prefix="aerospike-example", + ) as persister: + persister.initialize() + app = build_application(persister, args.app_id, args.partition_key) + if args.prompt: + *_, state = app.run(halt_after=["ai_response"], inputs={"prompt": args.prompt}) + print(state["response"]) + return + + while True: + prompt = input("you: ").strip() + if prompt.lower() in {"exit", "quit"}: + break + *_, state = app.run(halt_after=["ai_response"], inputs={"prompt": prompt}) + print("assistant:", state["response"]) + + +if __name__ == "__main__": + main() From b35d34556e79f1a0ecd85c3d751d3c46b6b3c372 Mon Sep 17 00:00:00 2001 From: PatersonProjects Date: Wed, 9 Sep 2026 09:38:18 -0700 Subject: [PATCH 4/4] Added health check to CI client, change name of class. --- .github/workflows/python-package.yml | 7 ++- burr/integrations/persisters/b_aerospike.py | 26 +++++++--- docs/reference/persister.rst | 4 +- examples/integrations/aerospike/README.md | 2 +- .../integrations/aerospike/aerospike_local.py | 6 +-- .../persisters/test_b_aerospike.py | 52 ++++++------------- 6 files changed, 45 insertions(+), 52 deletions(-) diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index d63a146d1..cb4e202c9 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -154,11 +154,16 @@ jobs: ports: - 6379:6379 aerospike: - image: aerospike/aerospike-server:latest + image: aerospike/aerospike-server:8.1.2.4 ports: - 3000:3000 - 3001:3001 - 3002:3002 + options: >- + --health-cmd "asinfo -v build" + --health-interval 5s + --health-timeout 3s + --health-retries 20 strategy: fail-fast: false matrix: diff --git a/burr/integrations/persisters/b_aerospike.py b/burr/integrations/persisters/b_aerospike.py index be025f71f..90ec4debe 100644 --- a/burr/integrations/persisters/b_aerospike.py +++ b/burr/integrations/persisters/b_aerospike.py @@ -56,6 +56,9 @@ _MAX_WRITE_ATTEMPTS = 3 _RETRYABLE_CODES = {9, -10, 7, 14} # Timeout, Connection, ClusterChange, KEY_BUSY +# Bounded wait for the head-set secondary index to become queryable. +_DEFAULT_INDEX_READY_TIMEOUT = 30.0 + class AerospikePersistenceError(Exception): """Base class for Aerospike persister errors.""" @@ -77,7 +80,7 @@ class AerospikePersistenceUncertainOutcomeError(AerospikePersistenceError): """Raised when a write retry budget is exhausted without a definitive result.""" -class AerospikePersister(persistence.BaseStatePersister): +class AerospikeBasePersister(persistence.BaseStatePersister): """Synchronous Aerospike-backed implementation of Burr's ``BaseStatePersister``. The persister stores one immutable history record per ``(partition_key, @@ -91,7 +94,7 @@ class AerospikePersister(persistence.BaseStatePersister): """ @classmethod - def from_config(cls, config: dict) -> "AerospikePersister": + def from_config(cls, config: dict) -> "AerospikeBasePersister": """Create a persister from a configuration dictionary.""" return cls.from_values(**config) @@ -107,7 +110,8 @@ def from_values( serde_kwargs: Optional[dict] = None, index_name: str = "burr_head_partition_idx", create_index: bool = True, - ) -> "AerospikePersister": + index_ready_timeout: float = _DEFAULT_INDEX_READY_TIMEOUT, + ) -> "AerospikeBasePersister": """Create a persister from seed hosts and client configuration. :param hosts: Aerospike seed hosts as ``[(host, port), ...]``. @@ -119,6 +123,8 @@ def from_values( :param serde_kwargs: Kwargs for Burr ``State`` serialization. :param index_name: Name of the secondary index on the head set. :param create_index: Whether ``initialize()`` may create the index. + :param index_ready_timeout: Seconds to wait for the head-set secondary + index to become queryable during ``initialize()``. """ if hosts is None: hosts = [("127.0.0.1", 3000)] @@ -137,6 +143,7 @@ def from_values( serde_kwargs=serde_kwargs, index_name=index_name, create_index=create_index, + index_ready_timeout=index_ready_timeout, _client_config=aerospike_config, _owned=True, ) @@ -152,6 +159,7 @@ def __init__( serde_kwargs: Optional[dict] = None, index_name: str = "burr_head_partition_idx", create_index: bool = True, + index_ready_timeout: float = _DEFAULT_INDEX_READY_TIMEOUT, _client_config: Optional[dict] = None, _owned: bool = False, ): @@ -170,6 +178,7 @@ def __init__( self.serde_kwargs = serde_kwargs or {} self.index_name = index_name self.create_index = create_index + self.index_ready_timeout = index_ready_timeout self._client_config = _client_config self._initialized = False self._closed = False @@ -193,11 +202,11 @@ def cleanup(self): def __getstate__(self) -> dict: if not self._owned: raise TypeError( - "An AerospikePersister constructed with an injected client cannot be pickled" + "An AerospikeBasePersister constructed with an injected client cannot be pickled" ) if self._client_config is None: raise TypeError( - "Cannot pickle an AerospikePersister without reconnectable client configuration" + "Cannot pickle an AerospikeBasePersister without reconnectable client configuration" ) state = self.__dict__.copy() del state["_client"] @@ -208,7 +217,7 @@ def __getstate__(self) -> dict: def __setstate__(self, state: dict): client_config = state.get("_client_config") if client_config is None: - raise TypeError("Cannot unpickle an AerospikePersister without client configuration") + raise TypeError("Cannot unpickle an AerospikeBasePersister without client configuration") self.__dict__.update(state) try: self._client = aerospike.client(client_config) @@ -243,7 +252,7 @@ def initialize(self): f"Failed to create secondary index '{self.index_name}': {e}" ) from e - deadline = time.monotonic() + 30.0 + deadline = time.monotonic() + self.index_ready_timeout attempt = 0 while True: try: @@ -251,7 +260,8 @@ def initialize(self): except (aerospike.exception.IndexNotFound, aerospike.exception.IndexNotReadable) as e: if time.monotonic() >= deadline: raise AerospikePersistenceInitializationError( - "Timed out waiting for the secondary index to become queryable" + f"Timed out waiting for the secondary index to become queryable " + f"after {self.index_ready_timeout:.0f}s" ) from e attempt += 1 self._backoff(attempt) diff --git a/docs/reference/persister.rst b/docs/reference/persister.rst index 3347f0b19..9f55008c4 100644 --- a/docs/reference/persister.rst +++ b/docs/reference/persister.rst @@ -42,7 +42,7 @@ We currently support the following database integrations: +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ | MongoDB | pymongo | :ref:`MongoDBBasePersister ` | ❌ | ❌ | +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ - | Aerospike | aerospike | :ref:`AerospikePersister ` | ❌ | ❌ | + | Aerospike | aerospike | :ref:`AerospikeBasePersister ` | ❌ | ❌ | +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ We follow the naming convention ``b_dependency-library``, where the ``b_`` is used to avoid name @@ -133,7 +133,7 @@ Currently we support the following, although we highly recommend you contribute .. _syncaerospikeref: -.. autoclass:: burr.integrations.persisters.b_aerospike.AerospikePersister +.. autoclass:: burr.integrations.persisters.b_aerospike.AerospikeBasePersister :members: .. automethod:: __init__ diff --git a/examples/integrations/aerospike/README.md b/examples/integrations/aerospike/README.md index 898622c4d..7b4d1aba4 100644 --- a/examples/integrations/aerospike/README.md +++ b/examples/integrations/aerospike/README.md @@ -42,7 +42,7 @@ Start a local Aerospike Community Edition database with Docker: ```bash docker run --name burr-aerospike -d \ -p 3000:3000 \ - aerospike/aerospike-server + aerospike/aerospike-server:8.1.2.4 ``` The example uses the default `test` namespace. Aerospike creates the `burr_state` and `burr_head` sets on the first write. diff --git a/examples/integrations/aerospike/aerospike_local.py b/examples/integrations/aerospike/aerospike_local.py index 02727010e..7eb6a83ba 100644 --- a/examples/integrations/aerospike/aerospike_local.py +++ b/examples/integrations/aerospike/aerospike_local.py @@ -20,7 +20,7 @@ import openai from burr.core import ApplicationBuilder, State, action -from burr.integrations.persisters.b_aerospike import AerospikePersister +from burr.integrations.persisters.b_aerospike import AerospikeBasePersister MODEL = "gpt-4o-mini" MAX_HISTORY_ITEMS = 20 @@ -45,7 +45,7 @@ def ai_response(state: State) -> State: return state.update(response=content, chat_history=chat_history) -def build_application(persister: AerospikePersister, app_id: str, partition_key: str): +def build_application(persister: AerospikeBasePersister, app_id: str, partition_key: str): return ( ApplicationBuilder() .with_actions(human_input, ai_response) @@ -72,7 +72,7 @@ def main(): parser.add_argument("--prompt") args = parser.parse_args() - with AerospikePersister.from_values( + with AerospikeBasePersister.from_values( hosts=[("127.0.0.1", 3000)], namespace="test", key_prefix="aerospike-example", diff --git a/tests/integrations/persisters/test_b_aerospike.py b/tests/integrations/persisters/test_b_aerospike.py index 20af9c71a..92082aa55 100644 --- a/tests/integrations/persisters/test_b_aerospike.py +++ b/tests/integrations/persisters/test_b_aerospike.py @@ -36,15 +36,14 @@ from burr.core import state from burr.core.persistence import BaseStatePersister from burr.integrations.persisters.b_aerospike import ( + AerospikeBasePersister, AerospikePersistenceInitializationError, - AerospikePersistenceSerializationError, - AerospikePersister, ) @pytest.fixture def aerospike_persister(): - persister = AerospikePersister.from_values(key_prefix=f"test-{uuid.uuid4().hex}") + persister = AerospikeBasePersister.from_values(key_prefix=f"test-{uuid.uuid4().hex}") persister.initialize() yield persister persister.cleanup() @@ -249,7 +248,7 @@ def test_repeated_initialization_is_idempotent(aerospike_persister): def test_validation_only_accepts_the_existing_compatible_index(aerospike_persister): - validator = AerospikePersister.from_values(create_index=False) + validator = AerospikeBasePersister.from_values(create_index=False) try: validator.initialize() assert validator.is_initialized() is True @@ -290,7 +289,7 @@ def operate(self, *args, **kwargs): def test_injected_client_constructs_a_synchronous_base_persister(): - persister = AerospikePersister(client=CallerOwnedClient()) + persister = AerospikeBasePersister(client=CallerOwnedClient()) assert isinstance(persister, BaseStatePersister) assert persister.is_async() is False @@ -298,7 +297,7 @@ def test_injected_client_constructs_a_synchronous_base_persister(): def test_cleanup_and_context_exit_never_close_an_injected_client(): client = CallerOwnedClient() - persister = AerospikePersister(client=client) + persister = AerospikeBasePersister(client=client) persister.cleanup() persister.cleanup() @@ -313,7 +312,7 @@ def test_cleanup_closes_an_internally_constructed_client_once(): with patch( "burr.integrations.persisters.b_aerospike.aerospike.client", return_value=client ) as factory: - persister = AerospikePersister.from_values() + persister = AerospikeBasePersister.from_values() persister.cleanup() persister.cleanup() @@ -336,7 +335,7 @@ def test_from_config_accepts_official_client_configuration(): with patch( "burr.integrations.persisters.b_aerospike.aerospike.client", return_value=client ) as factory: - persister = AerospikePersister.from_config(config) + persister = AerospikeBasePersister.from_config(config) try: supplied_config = factory.call_args.args[0] @@ -350,7 +349,7 @@ def test_from_config_accepts_official_client_configuration(): @pytest.mark.parametrize("sequence_id", [True, -(2**63) - 1, 2**63, 1.0, "1"]) def test_invalid_sequence_is_rejected_before_client_access(sequence_id): client = CallerOwnedClient() - persister = AerospikePersister(client=client) + persister = AerospikeBasePersister(client=client) with pytest.raises(ValueError, match="sequence"): persister.save( @@ -362,7 +361,7 @@ def test_invalid_sequence_is_rejected_before_client_access(sequence_id): def test_load_without_an_app_id_is_rejected_before_client_access(): client = CallerOwnedClient() - persister = AerospikePersister(client=client) + persister = AerospikeBasePersister(client=client) with pytest.raises(ValueError, match="app_id"): persister.load("pk", None) @@ -370,27 +369,6 @@ def test_load_without_an_app_id_is_rejected_before_client_access(): assert client.database_calls == 0 -def test_oversized_history_write_is_rejected_without_creating_or_advancing_head( - aerospike_persister, -): - """A record exceeding the namespace max-record-size must fail cleanly and leave no head.""" - # The test namespace uses max-record-size=1MB; a 2MB payload should exceed it. - large_payload = "x" * (2 * 1024 * 1024) - - with pytest.raises(AerospikePersistenceSerializationError, match="max-record-size|Oversized"): - aerospike_persister.save( - "pk", - "oversized-app", - 1, - "position", - state.State({"large": large_payload}), - "completed", - ) - - assert aerospike_persister.load("pk", "oversized-app") is None - assert aerospike_persister.list_app_ids("pk") == [] - - def test_concurrent_saves_to_same_application_advance_monotonically(aerospike_persister): """Concurrent saves to one application should create every checkpoint and leave the head at the max sequence.""" sequences = list(range(1, 11)) @@ -423,7 +401,7 @@ def test_owned_persister_uses_the_factory_connected_client(): with patch( "burr.integrations.persisters.b_aerospike.aerospike.client", return_value=client ): - persister = AerospikePersister.from_values() + persister = AerospikeBasePersister.from_values() try: client.connect.assert_not_called() @@ -451,7 +429,7 @@ def get(self, key, policy): def test_save_uses_write_policies_and_the_successful_operate_result(): client = RecordingClient() - persister = AerospikePersister(client=client) + persister = AerospikeBasePersister(client=client) persister.save( "partition", @@ -470,7 +448,7 @@ def test_save_uses_write_policies_and_the_successful_operate_result(): def test_initialize_fails_immediately_on_creation_error(): client = Mock() client.index_single_value_create.side_effect = aerospike.exception.AerospikeError() - persister = AerospikePersister(client=client) + persister = AerospikeBasePersister(client=client) with pytest.raises(AerospikePersistenceInitializationError, match="create"): persister.initialize() @@ -485,7 +463,7 @@ def test_initialize_retries_until_index_is_queryable(): aerospike.exception.IndexNotReadable(), [], ] - persister = AerospikePersister(client=client, create_index=False) + persister = AerospikeBasePersister(client=client, create_index=False) with patch.object(persister, "_backoff") as backoff: persister.initialize() @@ -498,7 +476,7 @@ def test_initialize_retries_until_index_is_queryable(): def test_initialize_times_out_when_index_never_becomes_queryable(): client = Mock() client.query.return_value.results.side_effect = aerospike.exception.IndexNotFound() - persister = AerospikePersister(client=client, create_index=False) + persister = AerospikeBasePersister(client=client, create_index=False) with patch( "burr.integrations.persisters.b_aerospike.time.monotonic", @@ -512,7 +490,7 @@ def test_initialize_times_out_when_index_never_becomes_queryable(): def test_initialize_fails_immediately_on_unexpected_query_error(): client = Mock() client.query.return_value.results.side_effect = aerospike.exception.AerospikeError() - persister = AerospikePersister(client=client, create_index=False) + persister = AerospikeBasePersister(client=client, create_index=False) with pytest.raises(AerospikePersistenceInitializationError, match="query readiness"): persister.initialize()