update the connector config after refreshing google calendar access token

This commit is contained in:
CREDO23 2025-08-26 11:49:31 +02:00
parent 9711af2b72
commit 85664f2ff8
2 changed files with 47 additions and 10 deletions

View file

@ -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)

View file

@ -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
)