diff --git a/surfsense_backend/app/connectors/google_calendar_connector.py b/surfsense_backend/app/connectors/google_calendar_connector.py index 3d7ca2e..38a7ae1 100644 --- a/surfsense_backend/app/connectors/google_calendar_connector.py +++ b/surfsense_backend/app/connectors/google_calendar_connector.py @@ -4,6 +4,7 @@ A module for retrieving calendar events from Google Calendar using Google OAuth Allows fetching events from specified calendars within date ranges using Google OAuth credentials. """ +import json from datetime import datetime from typing import Any @@ -12,6 +13,14 @@ from dateutil.parser import isoparse from google.auth.transport.requests import Request from google.oauth2.credentials import Credentials from googleapiclient.discovery import build +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select +from sqlalchemy.orm.attributes import flag_modified + +from app.db import ( + SearchSourceConnector, + SearchSourceConnectorType, +) class GoogleCalendarConnector: @@ -20,6 +29,8 @@ class GoogleCalendarConnector: def __init__( self, credentials: Credentials, + session: AsyncSession, + user_id: str, ): """ Initialize the GoogleCalendarConnector class. @@ -27,9 +38,13 @@ class GoogleCalendarConnector: credentials: Google OAuth Credentials object """ self._credentials = credentials + self._session = session + self._user_id = user_id self.service = None - def _get_credentials(self) -> Credentials: + async def _get_credentials( + self, + ) -> Credentials: """ Get valid Google OAuth credentials. Returns: @@ -60,12 +75,30 @@ class GoogleCalendarConnector: client_id=self._credentials.client_id, client_secret=self._credentials.client_secret, scopes=self._credentials.scopes, + expiry=self._credentials.expiry, ) # Refresh the token if needed - if self._credentials.expired or not self._credentials.valid: + if self._credentials.expired: try: self._credentials.refresh(Request()) + # Update the connector config in DB + if self._session: + result = await self._session.execute( + select(SearchSourceConnector).filter( + SearchSourceConnector.user_id == self._user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_CALENDAR_CONNECTOR, + ) + ) + connector = result.scalars().first() + if connector is None: + raise RuntimeError( + "GOOGLE_CALENDAR_CONNECTOR connector not found for current user; cannot persist refreshed token." + ) + connector.config = json.loads(self._credentials.to_json()) + flag_modified(connector, "config") + await self._session.commit() except Exception as e: raise Exception( f"Failed to refresh Google OAuth credentials: {e!s}" @@ -73,7 +106,7 @@ class GoogleCalendarConnector: return self._credentials - def _get_service(self): + async def _get_service(self): """ Get the Google Calendar service instance using Google OAuth credentials. Returns: @@ -86,20 +119,20 @@ class GoogleCalendarConnector: return self.service try: - credentials = self._get_credentials() + credentials = await self._get_credentials() self.service = build("calendar", "v3", credentials=credentials) return self.service except Exception as e: raise Exception(f"Failed to create Google Calendar service: {e!s}") from e - def get_calendars(self) -> tuple[list[dict[str, Any]], str | None]: + async def get_calendars(self) -> tuple[list[dict[str, Any]], str | None]: """ Fetch list of user's calendars using Google OAuth credentials. Returns: Tuple containing (calendars list, error message or None) """ try: - service = self._get_service() + service = await self._get_service() calendars_result = service.calendarList().list().execute() calendars = calendars_result.get("items", []) @@ -122,7 +155,7 @@ class GoogleCalendarConnector: except Exception as e: return [], f"Error fetching calendars: {e!s}" - def get_all_primary_calendar_events( + async def get_all_primary_calendar_events( self, start_date: str, end_date: str, @@ -136,7 +169,7 @@ class GoogleCalendarConnector: Tuple containing (events list, error message or None) """ try: - service = self._get_service() + service = await self._get_service() # Parse both dates dt_start = isoparse(start_date) diff --git a/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py b/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py index d3ac800..abc4925 100644 --- a/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py +++ b/surfsense_backend/app/tasks/connector_indexers/google_calendar_indexer.py @@ -81,6 +81,7 @@ async def index_google_calendar_events( return 0, f"Connector with ID {connector_id} not found" # Get the Google Calendar credentials from the connector config + exp = connector.config.get("expiry").replace("Z", "") credentials = Credentials( token=connector.config.get("token"), refresh_token=connector.config.get("refresh_token"), @@ -88,6 +89,7 @@ async def index_google_calendar_events( client_id=connector.config.get("client_id"), client_secret=connector.config.get("client_secret"), scopes=connector.config.get("scopes"), + expiry=datetime.fromisoformat(exp), ) if ( @@ -110,7 +112,9 @@ async def index_google_calendar_events( {"stage": "client_initialization"}, ) - calendar_client = GoogleCalendarConnector(credentials=credentials) + calendar_client = GoogleCalendarConnector( + credentials=credentials, session=session, user_id=user_id + ) # Calculate date range if start_date is None or end_date is None: @@ -169,7 +173,7 @@ async def index_google_calendar_events( # Get events within date range from primary calendar try: - events, error = calendar_client.get_all_primary_calendar_events( + events, error = await calendar_client.get_all_primary_calendar_events( start_date=start_date_str, end_date=end_date_str )