diff --git a/.gitignore b/.gitignore index a92fd35b4..9c29d3c3d 100644 --- a/.gitignore +++ b/.gitignore @@ -193,3 +193,4 @@ burr/tracking/server/build examples/*/statemachine examples/*/*/statemachine .vscode +google_creds.json diff --git a/burr/integrations/persisters/b_google_sheets.py b/burr/integrations/persisters/b_google_sheets.py new file mode 100644 index 000000000..5daab5f64 --- /dev/null +++ b/burr/integrations/persisters/b_google_sheets.py @@ -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 \ No newline at end of file diff --git a/test_sheets.py b/test_sheets.py new file mode 100644 index 000000000..1bbf50f9d --- /dev/null +++ b/test_sheets.py @@ -0,0 +1,15 @@ +from burr.integrations.persisters.b_google_sheets import GoogleSheetsPersister + +persister = GoogleSheetsPersister( + spreadsheet_id="1ur4GqM0tQQyCRn3lAj5bZk3DAeZBvrbYRwKok780hHc", + credentials_file="google_creds.json" +) + +# 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()) \ No newline at end of file