Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -193,3 +193,4 @@ burr/tracking/server/build
examples/*/statemachine
examples/*/*/statemachine
.vscode
google_creds.json
78 changes: 78 additions & 0 deletions burr/integrations/persisters/b_google_sheets.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
import json
from burr.core import persistence
from googleapiclient.discovery import build
from google.oauth2.service_account import Credentials

class GoogleSheetsPersister(persistence.BaseStatePersister):
"""
Persister that stores Burr state in Google Sheets.
"""

def __init__(self, spreadsheet_id: str, credentials_file: str):
self.spreadsheet_id = spreadsheet_id
scopes = ["https://www.googleapis.com/auth/spreadsheets"]
creds = Credentials.from_service_account_file(
credentials_file,
scopes=scopes)
self.service = build("sheets", "v4", credentials=creds)



def save(self, app_id: str, partition_key: str, state: dict):
state_json = json.dumps(state)

values = [[app_id, partition_key, state_json]]
body = {"values": values}

self.service.spreadsheets().values().append(
spreadsheetId=self.spreadsheet_id,
range="states!A:C",
valueInputOption="RAW",
body=body).execute()



def load(self, app_id: str, partition_key: str):
result = self.service.spreadsheets().values().get(
spreadsheetId=self.spreadsheet_id,
range="states!A:C").execute()

rows = result.get("values", [])

for row in rows:
if len(row) >= 3 and row[0] == app_id and row[1] == partition_key:
return json.loads(row[2])

return None



def list(self, app_id: str):
result = self.service.spreadsheets().values().get(
spreadsheetId=self.spreadsheet_id,
range="states!A:C").execute()

rows = result.get("values", [])
partition_keys = []

for row in rows:
if len(row) >= 2 and row[0] == app_id:
partition_keys.append(row[1])

return partition_keys



def list_app_ids(self):
result = self.service.spreadsheets().values().get(
spreadsheetId=self.spreadsheet_id,
range="states!A:C").execute()

rows = result.get("values", [])
app_ids = []

for row in rows[1:]:
if len(row) >= 1 and row[0] not in app_ids:
app_ids.append(row[0])

return app_ids
15 changes: 15 additions & 0 deletions test_sheets.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from burr.integrations.persisters.b_google_sheets import GoogleSheetsPersister

persister = GoogleSheetsPersister(
spreadsheet_id="1ur4GqM0tQQyCRn3lAj5bZk3DAeZBvrbYRwKok780hHc",
credentials_file="google_creds.json"
)
Comment on lines +3 to +6
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah how would we implement a test for this?


# Save some data
persister.save("demo_app", "user1", {"hello": "world"})

# Load the data
data = persister.load("demo_app", "user1")

print("Loaded data:", data)
print("App IDs:", persister.list_app_ids())