Merge pull request #407 from chetan/fix/service-mem-leak
fix/service mem leak
This commit is contained in:
@@ -784,6 +784,9 @@ def get_user_info(credentials: Credentials) -> Optional[Dict[str, Any]]:
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error fetching user info: {e}")
|
||||
return None
|
||||
finally:
|
||||
if service:
|
||||
service.close()
|
||||
|
||||
|
||||
# --- Centralized Google Service Authentication ---
|
||||
|
||||
@@ -15,6 +15,7 @@ from dataclasses import dataclass
|
||||
|
||||
from fastmcp.server.auth import AccessToken
|
||||
from google.oauth2.credentials import Credentials
|
||||
from auth.oauth_config import is_external_oauth21_provider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -743,7 +744,8 @@ def ensure_session_from_access_token(
|
||||
else:
|
||||
store_expiry = credentials.expiry
|
||||
|
||||
if email:
|
||||
# Skip session storage for external OAuth 2.1 to prevent memory leak from ephemeral tokens
|
||||
if email and not is_external_oauth21_provider():
|
||||
try:
|
||||
store = get_oauth21_session_store()
|
||||
store.store_session(
|
||||
|
||||
@@ -4,6 +4,7 @@ import logging
|
||||
import re
|
||||
from functools import wraps
|
||||
from typing import Dict, List, Optional, Any, Callable, Union, Tuple
|
||||
from contextlib import ExitStack
|
||||
|
||||
from google.auth.exceptions import RefreshError
|
||||
from googleapiclient.discovery import build
|
||||
@@ -687,6 +688,9 @@ def require_google_service(
|
||||
e, actual_user_email, service_name
|
||||
)
|
||||
raise GoogleAuthenticationError(error_message)
|
||||
finally:
|
||||
if service:
|
||||
service.close()
|
||||
|
||||
# Set the wrapper's signature to the one without 'service'
|
||||
wrapper.__signature__ = wrapper_sig
|
||||
@@ -757,6 +761,7 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
|
||||
)
|
||||
|
||||
# Authenticate all services
|
||||
with ExitStack() as stack:
|
||||
for config in service_configs:
|
||||
service_type = config["service_type"]
|
||||
scopes = config["scopes"]
|
||||
@@ -804,6 +809,7 @@ def require_multiple_services(service_configs: List[Dict[str, Any]]):
|
||||
|
||||
# Inject service with specified parameter name
|
||||
kwargs[param_name] = service
|
||||
stack.callback(service.close)
|
||||
|
||||
except GoogleAuthenticationError as e:
|
||||
logger.error(
|
||||
|
||||
@@ -671,9 +671,14 @@ async def create_event(
|
||||
event_body["attachments"] = []
|
||||
drive_service = None
|
||||
try:
|
||||
drive_service = service._http and build("drive", "v3", http=service._http)
|
||||
try:
|
||||
drive_service = service._http and build(
|
||||
"drive", "v3", http=service._http
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not build Drive service for MIME type lookup: {e}")
|
||||
logger.warning(
|
||||
f"Could not build Drive service for MIME type lookup: {e}"
|
||||
)
|
||||
for att in attachments:
|
||||
file_id = None
|
||||
if att.startswith("https://"):
|
||||
@@ -726,6 +731,9 @@ async def create_event(
|
||||
"mimeType": mime_type,
|
||||
}
|
||||
)
|
||||
finally:
|
||||
if drive_service:
|
||||
drive_service.close()
|
||||
created_event = await asyncio.to_thread(
|
||||
lambda: service.events()
|
||||
.insert(
|
||||
|
||||
Reference in New Issue
Block a user