From d01a7882246b0383ad3a853200a7bcb252346bf5 Mon Sep 17 00:00:00 2001 From: prvn0x Date: Sat, 12 Sep 2026 19:47:27 +0530 Subject: [PATCH] Add Google Sheets state persister Closes #213 Co-authored-by: goutamk09 --- .../persisters/b_google_sheets.py | 246 ++++++++++++++++++ docs/reference/persister.rst | 10 + pyproject.toml | 5 + .../persisters/test_b_google_sheets.py | 156 +++++++++++ 4 files changed, 417 insertions(+) create mode 100644 burr/integrations/persisters/b_google_sheets.py create mode 100644 tests/integrations/persisters/test_b_google_sheets.py diff --git a/burr/integrations/persisters/b_google_sheets.py b/burr/integrations/persisters/b_google_sheets.py new file mode 100644 index 000000000..758f8ebd1 --- /dev/null +++ b/burr/integrations/persisters/b_google_sheets.py @@ -0,0 +1,246 @@ +# 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. + +from burr.integrations import base + +try: + import gspread +except ImportError as e: + base.require_plugin(e, "google-sheets") + +import json +import logging +from datetime import datetime, timezone +from typing import Any, Literal, Optional + +from burr.core import persistence, state + +logger = logging.getLogger(__name__) + +COLUMNS = [ + "partition_key", + "app_id", + "sequence_id", + "position", + "state", + "created_at", + "status", +] + +_PARTITION_KEY = 0 +_APP_ID = 1 +_SEQUENCE_ID = 2 +_POSITION = 3 +_STATE = 4 +_CREATED_AT = 5 +_STATUS = 6 + + +def _encode_partition_key(partition_key: Optional[str]) -> str: + """Google Sheets has no concept of null, so ``None`` is stored as an empty cell.""" + return "" if partition_key is None else partition_key + + +class GoogleSheetsBasePersister(persistence.BaseStatePersister): + """Persister that stores state in a Google Sheets worksheet. + + .. warning:: + This persister is in **beta**. Google Sheets is not a database -- the Sheets API + enforces per-user quotas (on the order of 60 reads and 300 writes per minute), every + value is stored as a string, and there are no indexes, so reads scan the whole + worksheet. It is intended for demos, prototypes, and low-volume applications where + a human-readable, shareable store is worth more than throughput. + + Each saved step is appended as a single row. The first row of the worksheet is a header + (see :py:data:`COLUMNS`), written by :py:meth:`initialize`, and is skipped by every read. + + Use this class directly if you want to control the ``gspread`` worksheet object yourself, + or use :py:meth:`from_values` to build one from a service account file. + """ + + @classmethod + def from_config(cls, config: dict) -> "GoogleSheetsBasePersister": + """Creates a new instance of the GoogleSheetsBasePersister from a configuration dictionary.""" + return cls.from_values(**config) + + @classmethod + def from_values( + cls, + service_account_file: str, + spreadsheet_key: str, + worksheet_name: str = "Sheet1", + serde_kwargs: dict = None, + ) -> "GoogleSheetsBasePersister": + """Creates a new instance of the GoogleSheetsBasePersister from passed in values. + + :param service_account_file: path to a Google service account JSON key file. The + service account must have edit access to the target spreadsheet. + :param spreadsheet_key: the spreadsheet ID, i.e. the ```` in + ``https://docs.google.com/spreadsheets/d//edit``. + :param worksheet_name: the worksheet (tab) within the spreadsheet to use. + :param serde_kwargs: serialization and deserialization keyword arguments to pass to state SERDE. + """ + client = gspread.service_account(filename=service_account_file) + worksheet = client.open_by_key(spreadsheet_key).worksheet(worksheet_name) + return cls(worksheet, serde_kwargs) + + def __init__(self, worksheet: Any, serde_kwargs: dict = None): + """Initializes the GoogleSheetsBasePersister class. + + :param worksheet: a ``gspread`` worksheet object to read and write. + :param serde_kwargs: serialization and deserialization keyword arguments to pass to state SERDE. + """ + self.worksheet = worksheet + self.serde_kwargs = serde_kwargs or {} + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def set_serde_kwargs(self, serde_kwargs: dict): + """Sets the serde_kwargs for the persister.""" + self.serde_kwargs = serde_kwargs + + def initialize(self): + """Writes the header row if the worksheet does not already have one. + + This is idempotent -- calling it on an initialized worksheet is a no-op. + """ + if self.is_initialized(): + return + self.worksheet.append_row(COLUMNS) + + def is_initialized(self) -> bool: + """Whether the worksheet has the expected header row.""" + rows = self.worksheet.get_all_values() + return bool(rows) and rows[0] == COLUMNS + + def _data_rows(self) -> list[list[str]]: + """Returns the worksheet rows with the header stripped. + + Every read goes through here so that the header is skipped consistently. + """ + rows = self.worksheet.get_all_values() + if rows and rows[0] == COLUMNS: + return rows[1:] + return rows + + def list_app_ids(self, partition_key: str, **kwargs) -> list[str]: + """List the app ids for a given partition key, most recently written first.""" + encoded = _encode_partition_key(partition_key) + last_written = {} + for index, row in enumerate(self._data_rows()): + if len(row) < len(COLUMNS) or row[_PARTITION_KEY] != encoded: + continue + last_written[row[_APP_ID]] = index + return [ + app_id + for app_id, _ in sorted(last_written.items(), key=lambda item: item[1], reverse=True) + ] + + def load( + self, partition_key: str, app_id: str, sequence_id: int = None, **kwargs + ) -> Optional[persistence.PersistedStateData]: + """Load the state data for a given partition key, app id, and sequence id. + + If ``sequence_id`` is not given, the state of the most recently written row for this + app is returned. Because rows are only ever appended, the *last* matching row is + always the current one -- earlier rows are prior steps. + + :param partition_key: + :param app_id: + :param sequence_id: + :param kwargs: + :return: PersistedStateData or None if no matching row exists. + """ + encoded = _encode_partition_key(partition_key) + match = None + for row in self._data_rows(): + if len(row) < len(COLUMNS): + continue + if row[_PARTITION_KEY] != encoded or row[_APP_ID] != app_id: + continue + if sequence_id is not None and row[_SEQUENCE_ID] != str(sequence_id): + continue + match = row # keep scanning; the last match wins + if match is None: + return None + _state = state.State.deserialize(json.loads(match[_STATE]), **self.serde_kwargs) + return { + "partition_key": partition_key, + "app_id": match[_APP_ID], + "sequence_id": int(match[_SEQUENCE_ID]), + "position": match[_POSITION], + "state": _state, + "created_at": match[_CREATED_AT], + "status": match[_STATUS], + } + + def save( + self, + partition_key: Optional[str], + app_id: str, + sequence_id: int, + position: str, + state: state.State, + status: Literal["completed", "failed"], + **kwargs, + ): + """Append the state data as a new row in the worksheet. + + :param partition_key: + :param app_id: + :param sequence_id: + :param position: + :param state: + :param status: + :param kwargs: + :raises ValueError: if a row for this (partition_key, app_id, sequence_id, position) + already exists. + """ + encoded = _encode_partition_key(partition_key) + for row in self._data_rows(): + if len(row) < len(COLUMNS): + continue + if ( + row[_PARTITION_KEY] == encoded + and row[_APP_ID] == app_id + and row[_SEQUENCE_ID] == str(sequence_id) + and row[_POSITION] == position + ): + raise ValueError( + f"partition_key:app_id:sequence_id:position" + f"[{encoded}:{app_id}:{sequence_id}:{position}] already exists." + ) + json_state = json.dumps(state.serialize(**self.serde_kwargs)) + self.worksheet.append_row( + [ + encoded, + app_id, + str(sequence_id), + position, + json_state, + datetime.now(timezone.utc).isoformat(), + status, + ] + ) + + def cleanup(self): + """No-op -- the Google Sheets client holds no long-lived connection to close.""" + pass diff --git a/docs/reference/persister.rst b/docs/reference/persister.rst index 9564bc81e..90525ead5 100644 --- a/docs/reference/persister.rst +++ b/docs/reference/persister.rst @@ -42,6 +42,9 @@ We currently support the following database integrations: +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ | MongoDB | pymongo | :ref:`MongoDBBasePersister ` | ❌ | ❌ | +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ + | Google | gspread | :ref:`GoogleSheetsBasePersister ` | ❌ | ❌ | + | Sheets | | | | | + +-------------+-----------+-----------------------------------------------------+---------------+-----------------------------------------------------+ 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 +132,13 @@ Currently we support the following, although we highly recommend you contribute .. automethod:: __init__ +.. _syncsheetsref: + +.. autoclass:: burr.integrations.persisters.b_google_sheets.GoogleSheetsBasePersister + :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 7fba971cc..c2bc3e716 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,10 @@ redis = [ "redis" ] +google-sheets = [ + "gspread" +] + release = [ "jinja2", ] @@ -102,6 +106,7 @@ tests = [ "apache-burr[psycopg2]", "apache-burr[pymongo]", "apache-burr[redis]", + "apache-burr[google-sheets]", "apache-burr[opentelemetry]", "apache-burr[langfuse]", "apache-burr[haystack]", diff --git a/tests/integrations/persisters/test_b_google_sheets.py b/tests/integrations/persisters/test_b_google_sheets.py new file mode 100644 index 000000000..c7553d607 --- /dev/null +++ b/tests/integrations/persisters/test_b_google_sheets.py @@ -0,0 +1,156 @@ +# 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 pytest + +from burr.core import state +from burr.integrations.persisters.b_google_sheets import COLUMNS, GoogleSheetsBasePersister + + +class FakeWorksheet: + """In-memory stand-in for a ``gspread`` worksheet. + + Only the two methods the persister uses are implemented. Keeping the persister's + dependency surface this small is what makes it testable without credentials or network. + Google Sheets returns every cell as a string, so this does too. + """ + + def __init__(self, rows: list[list[str]] = None): + self.rows = [list(row) for row in (rows or [])] + + def get_all_values(self) -> list[list[str]]: + return [list(row) for row in self.rows] + + def append_row(self, values: list) -> None: + self.rows.append([str(value) for value in values]) + + +@pytest.fixture +def persister(): + p = GoogleSheetsBasePersister(FakeWorksheet()) + p.initialize() + return p + + +def test_initialize_writes_header(persister): + assert persister.worksheet.rows[0] == COLUMNS + assert persister.is_initialized() + + +def test_initialize_is_idempotent(persister): + persister.initialize() + persister.initialize() + assert persister.worksheet.rows.count(COLUMNS) == 1 + + +def test_is_initialized_false_on_empty_worksheet(): + assert not GoogleSheetsBasePersister(FakeWorksheet()).is_initialized() + + +def test_save_and_load_state(persister): + persister.save("pk", "app_id", 1, "pos", state.State({"a": 1, "b": 2}), "completed") + data = persister.load("pk", "app_id", 1) + assert data["state"].get_all() == {"a": 1, "b": 2} + assert data["sequence_id"] == 1 + assert data["position"] == "pos" + assert data["status"] == "completed" + + +def test_load_without_sequence_id_returns_latest(persister): + """Rows are append-only, so the last matching row is the current state.""" + persister.save("pk", "app_id", 1, "first", state.State({"count": 1}), "completed") + persister.save("pk", "app_id", 2, "second", state.State({"count": 2}), "completed") + persister.save("pk", "app_id", 3, "third", state.State({"count": 3}), "completed") + data = persister.load("pk", "app_id") + assert data["sequence_id"] == 3 + assert data["position"] == "third" + assert data["state"].get_all() == {"count": 3} + + +def test_load_specific_sequence_id(persister): + persister.save("pk", "app_id", 1, "first", state.State({"count": 1}), "completed") + persister.save("pk", "app_id", 2, "second", state.State({"count": 2}), "completed") + data = persister.load("pk", "app_id", 1) + assert data["sequence_id"] == 1 + assert data["state"].get_all() == {"count": 1} + + +def test_load_returns_none_when_missing(persister): + assert persister.load("pk", "does-not-exist") is None + + +def test_load_is_scoped_to_partition_key(persister): + persister.save("pk1", "app_id", 1, "pos", state.State({"a": 1}), "completed") + persister.save("pk2", "app_id", 1, "pos", state.State({"a": 2}), "completed") + assert persister.load("pk1", "app_id")["state"].get_all() == {"a": 1} + assert persister.load("pk2", "app_id")["state"].get_all() == {"a": 2} + + +def test_list_app_ids_skips_header_and_dedupes(persister): + persister.save("pk", "app_id1", 1, "pos", state.State({"a": 1}), "completed") + persister.save("pk", "app_id2", 1, "pos", state.State({"b": 2}), "completed") + persister.save("pk", "app_id1", 2, "pos", state.State({"a": 3}), "completed") + app_ids = persister.list_app_ids("pk") + assert COLUMNS[1] not in app_ids, "header row must never be treated as data" + # app_id1 was written most recently, so it comes first + assert app_ids == ["app_id1", "app_id2"] + + +def test_list_app_ids_is_scoped_to_partition_key(persister): + persister.save("pk1", "app_id1", 1, "pos", state.State({"a": 1}), "completed") + persister.save("pk2", "app_id2", 1, "pos", state.State({"b": 2}), "completed") + assert persister.list_app_ids("pk1") == ["app_id1"] + assert persister.list_app_ids("pk2") == ["app_id2"] + + +def test_list_app_ids_empty_for_unknown_partition_key(persister): + assert persister.list_app_ids("nope") == [] + + +def test_save_duplicate_key_raises(persister): + persister.save("pk", "app_id", 1, "pos", state.State({"a": 1}), "completed") + with pytest.raises(ValueError, match="already exists"): + persister.save("pk", "app_id", 1, "pos", state.State({"a": 2}), "completed") + + +def test_same_sequence_id_different_position_is_allowed(persister): + persister.save("pk", "app_id", 1, "pos_a", state.State({"a": 1}), "completed") + persister.save("pk", "app_id", 1, "pos_b", state.State({"a": 2}), "completed") + assert persister.load("pk", "app_id", 1)["position"] == "pos_b" + + +def test_none_partition_key_roundtrip(persister): + persister.save(None, "app_id", 1, "pos", state.State({"a": 1}), "completed") + data = persister.load(None, "app_id") + assert data is not None + assert data["state"].get_all() == {"a": 1} + assert persister.list_app_ids(None) == ["app_id"] + + +def test_failed_status_is_persisted(persister): + persister.save("pk", "app_id", 1, "pos", state.State({"a": 1}), "failed") + assert persister.load("pk", "app_id")["status"] == "failed" + + +def test_reads_worksheet_without_header(persister): + """A worksheet written by an older version has no header row; reads must not drop a row.""" + raw = FakeWorksheet() + p = GoogleSheetsBasePersister(raw) + p.save("pk", "app_id", 1, "pos", state.State({"a": 1}), "completed") + assert not p.is_initialized() + assert p.load("pk", "app_id")["state"].get_all() == {"a": 1} + assert p.list_app_ids("pk") == ["app_id"]