mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-05 11:27:39 +00:00
df16a80f25
### What problem does this PR solve? This feature is primarily ported from the [Onyx](https://github.com/onyx-dot-app/onyx) project with necessary modifications. Thanks for such a brilliant project. Minor: consistently use `google_drive` rather than `google_driver`. <img width="566" height="731" alt="image" src="https://github.com/user-attachments/assets/6f64e70e-881e-42c7-b45f-809d3e0024a4" /> <img width="904" height="830" alt="image" src="https://github.com/user-attachments/assets/dfa7d1ef-819a-4a82-8c52-0999f48ed4a6" /> <img width="911" height="869" alt="image" src="https://github.com/user-attachments/assets/39e792fb-9fbe-4f3d-9b3c-b2265186bc22" /> <img width="947" height="323" alt="image" src="https://github.com/user-attachments/assets/27d70e96-d9c0-42d9-8c89-276919b6d61d" /> ### Type of change - [x] New Feature (non-breaking change which adds functionality)
121 lines
3.9 KiB
Python
121 lines
3.9 KiB
Python
import logging
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from google.auth.exceptions import RefreshError # type: ignore
|
|
from google.oauth2.credentials import Credentials as OAuthCredentials # type: ignore # type: ignore
|
|
from google.oauth2.service_account import Credentials as ServiceAccountCredentials # type: ignore # type: ignore
|
|
from googleapiclient.discovery import (
|
|
Resource, # type: ignore
|
|
build, # type: ignore
|
|
)
|
|
|
|
|
|
class GoogleDriveService(Resource):
|
|
pass
|
|
|
|
|
|
class GoogleDocsService(Resource):
|
|
pass
|
|
|
|
|
|
class AdminService(Resource):
|
|
pass
|
|
|
|
|
|
class GmailService(Resource):
|
|
pass
|
|
|
|
|
|
class RefreshableDriveObject:
|
|
"""
|
|
Running Google drive service retrieval functions
|
|
involves accessing methods of the service object (ie. files().list())
|
|
which can raise a RefreshError if the access token is expired.
|
|
This class is a wrapper that propagates the ability to refresh the access token
|
|
and retry the final retrieval function until execute() is called.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
call_stack: Callable[[ServiceAccountCredentials | OAuthCredentials], Any],
|
|
creds: ServiceAccountCredentials | OAuthCredentials,
|
|
creds_getter: Callable[..., ServiceAccountCredentials | OAuthCredentials],
|
|
):
|
|
self.call_stack = call_stack
|
|
self.creds = creds
|
|
self.creds_getter = creds_getter
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
if name == "execute":
|
|
return self.make_refreshable_execute()
|
|
return RefreshableDriveObject(
|
|
lambda creds: getattr(self.call_stack(creds), name),
|
|
self.creds,
|
|
self.creds_getter,
|
|
)
|
|
|
|
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
|
return RefreshableDriveObject(
|
|
lambda creds: self.call_stack(creds)(*args, **kwargs),
|
|
self.creds,
|
|
self.creds_getter,
|
|
)
|
|
|
|
def make_refreshable_execute(self) -> Callable:
|
|
def execute(*args: Any, **kwargs: Any) -> Any:
|
|
try:
|
|
return self.call_stack(self.creds).execute(*args, **kwargs)
|
|
except RefreshError as e:
|
|
logging.warning(f"RefreshError, going to attempt a creds refresh and retry: {e}")
|
|
# Refresh the access token
|
|
self.creds = self.creds_getter()
|
|
return self.call_stack(self.creds).execute(*args, **kwargs)
|
|
|
|
return execute
|
|
|
|
|
|
def _get_google_service(
|
|
service_name: str,
|
|
service_version: str,
|
|
creds: ServiceAccountCredentials | OAuthCredentials,
|
|
user_email: str | None = None,
|
|
) -> GoogleDriveService | GoogleDocsService | AdminService | GmailService:
|
|
service: Resource
|
|
if isinstance(creds, ServiceAccountCredentials):
|
|
# NOTE: https://developers.google.com/identity/protocols/oauth2/service-account#error-codes
|
|
creds = creds.with_subject(user_email)
|
|
service = build(service_name, service_version, credentials=creds)
|
|
elif isinstance(creds, OAuthCredentials):
|
|
service = build(service_name, service_version, credentials=creds)
|
|
|
|
return service
|
|
|
|
|
|
def get_google_docs_service(
|
|
creds: ServiceAccountCredentials | OAuthCredentials,
|
|
user_email: str | None = None,
|
|
) -> GoogleDocsService:
|
|
return _get_google_service("docs", "v1", creds, user_email)
|
|
|
|
|
|
def get_drive_service(
|
|
creds: ServiceAccountCredentials | OAuthCredentials,
|
|
user_email: str | None = None,
|
|
) -> GoogleDriveService:
|
|
return _get_google_service("drive", "v3", creds, user_email)
|
|
|
|
|
|
def get_admin_service(
|
|
creds: ServiceAccountCredentials | OAuthCredentials,
|
|
user_email: str | None = None,
|
|
) -> AdminService:
|
|
return _get_google_service("admin", "directory_v1", creds, user_email)
|
|
|
|
|
|
def get_gmail_service(
|
|
creds: ServiceAccountCredentials | OAuthCredentials,
|
|
user_email: str | None = None,
|
|
) -> GmailService:
|
|
return _get_google_service("gmail", "v1", creds, user_email)
|