mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-18 11:51:41 +03:00
d41a0b4be4
# Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate):
132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
import json
|
|
import os
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
from google_auth_oauthlib.flow import Flow
|
|
from logger import get_logger
|
|
from middlewares.auth import AuthBearer, get_current_user
|
|
from modules.sync.dto.inputs import SyncsUserInput, SyncUserUpdateInput
|
|
from modules.sync.service.sync_service import SyncService, SyncUserService
|
|
from modules.user.entity.user_identity import UserIdentity
|
|
|
|
# Set environment variable for OAuthlib
|
|
os.environ["OAUTHLIB_INSECURE_TRANSPORT"] = "1"
|
|
|
|
# Initialize logger
|
|
logger = get_logger(__name__)
|
|
|
|
# Initialize sync service
|
|
sync_service = SyncService()
|
|
sync_user_service = SyncUserService()
|
|
|
|
# Initialize API router
|
|
google_sync_router = APIRouter()
|
|
|
|
# Constants
|
|
SCOPES = [
|
|
"https://www.googleapis.com/auth/drive.metadata.readonly",
|
|
"https://www.googleapis.com/auth/drive.readonly",
|
|
]
|
|
BACKEND_URL = os.getenv("BACKEND_URL", "http://localhost:5050")
|
|
BASE_REDIRECT_URI = f"{BACKEND_URL}/sync/google/oauth2callback"
|
|
|
|
# Create credentials content from environment variables
|
|
CLIENT_SECRETS_FILE_CONTENT = {
|
|
"installed": {
|
|
"client_id": os.getenv("GOOGLE_CLIENT_ID"),
|
|
"project_id": os.getenv("GOOGLE_PROJECT_ID"),
|
|
"auth_uri": os.getenv("GOOGLE_AUTH_URI"),
|
|
"token_uri": os.getenv("GOOGLE_TOKEN_URI"),
|
|
"auth_provider_x509_cert_url": os.getenv("GOOGLE_AUTH_PROVIDER_CERT_URL"),
|
|
"client_secret": os.getenv("GOOGLE_CLIENT_SECRET"),
|
|
"redirect_uris": [os.getenv("GOOGLE_REDIRECT_URI")],
|
|
}
|
|
}
|
|
|
|
|
|
@google_sync_router.get(
|
|
"/sync/google/authorize",
|
|
dependencies=[Depends(AuthBearer())],
|
|
tags=["Sync"],
|
|
)
|
|
def authorize_google(
|
|
request: Request, name: str, current_user: UserIdentity = Depends(get_current_user)
|
|
):
|
|
"""
|
|
Authorize Google Drive sync for the current user.
|
|
|
|
Args:
|
|
request (Request): The request object.
|
|
current_user (UserIdentity): The current authenticated user.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the authorization URL.
|
|
"""
|
|
logger.debug(f"Authorizing Google Drive sync for user: {current_user.id}")
|
|
redirect_uri = f"{BASE_REDIRECT_URI}?user_id={current_user.id}"
|
|
flow = Flow.from_client_config(
|
|
CLIENT_SECRETS_FILE_CONTENT, scopes=SCOPES, redirect_uri=redirect_uri
|
|
)
|
|
authorization_url, state = flow.authorization_url(
|
|
access_type="offline", include_granted_scopes="true"
|
|
)
|
|
logger.info(
|
|
f"Generated authorization URL: {authorization_url} for user: {current_user.id}"
|
|
)
|
|
sync_user_input = SyncsUserInput(
|
|
name=name,
|
|
user_id=str(current_user.id),
|
|
provider="Google",
|
|
credentials={},
|
|
state={"state": state},
|
|
)
|
|
sync_user_service.create_sync_user(sync_user_input)
|
|
return {"authorization_url": authorization_url}
|
|
|
|
|
|
@google_sync_router.get("/sync/google/oauth2callback", tags=["Sync"])
|
|
def oauth2callback_google(request: Request):
|
|
"""
|
|
Handle OAuth2 callback from Google.
|
|
|
|
Args:
|
|
request (Request): The request object.
|
|
|
|
Returns:
|
|
dict: A dictionary containing a success message.
|
|
"""
|
|
state = request.query_params.get("state")
|
|
state_dict = {"state": state}
|
|
current_user = request.query_params.get("user_id")
|
|
logger.debug(
|
|
f"Handling OAuth2 callback for user: {current_user} with state: {state}"
|
|
)
|
|
sync_user_state = sync_user_service.get_sync_user_by_state(state_dict)
|
|
logger.info(f"Retrieved sync user state: {sync_user_state}")
|
|
|
|
if state_dict != sync_user_state["state"]:
|
|
logger.error("Invalid state parameter")
|
|
raise HTTPException(status_code=400, detail="Invalid state parameter")
|
|
if sync_user_state.get("user_id") != current_user:
|
|
logger.error("Invalid user")
|
|
raise HTTPException(status_code=400, detail="Invalid user")
|
|
|
|
redirect_uri = f"{BASE_REDIRECT_URI}?user_id={current_user}"
|
|
flow = Flow.from_client_config(
|
|
CLIENT_SECRETS_FILE_CONTENT,
|
|
scopes=SCOPES,
|
|
state=state,
|
|
redirect_uri=redirect_uri,
|
|
)
|
|
flow.fetch_token(authorization_response=str(request.url))
|
|
creds = flow.credentials
|
|
logger.info(f"Fetched OAuth2 token for user: {current_user}")
|
|
|
|
sync_user_input = SyncUserUpdateInput(
|
|
credentials=json.loads(creds.to_json()),
|
|
state={},
|
|
)
|
|
sync_user_service.update_sync_user(current_user, state_dict, sync_user_input)
|
|
logger.info(f"Google Drive sync created successfully for user: {current_user}")
|
|
return {"message": "Google Drive sync created successfully"}
|