From 2220314041349502767ea5e7bfa144d9828b23d4 Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Fri, 20 Feb 2026 16:05:54 +0200 Subject: [PATCH] add google drive tool metadata service --- .../app/services/google_drive/__init__.py | 11 ++ .../google_drive/tool_metadata_service.py | 147 ++++++++++++++++++ 2 files changed, 158 insertions(+) create mode 100644 surfsense_backend/app/services/google_drive/__init__.py create mode 100644 surfsense_backend/app/services/google_drive/tool_metadata_service.py diff --git a/surfsense_backend/app/services/google_drive/__init__.py b/surfsense_backend/app/services/google_drive/__init__.py new file mode 100644 index 000000000..5958a1761 --- /dev/null +++ b/surfsense_backend/app/services/google_drive/__init__.py @@ -0,0 +1,11 @@ +from app.services.google_drive.tool_metadata_service import ( + GoogleDriveAccount, + GoogleDriveFile, + GoogleDriveToolMetadataService, +) + +__all__ = [ + "GoogleDriveAccount", + "GoogleDriveFile", + "GoogleDriveToolMetadataService", +] diff --git a/surfsense_backend/app/services/google_drive/tool_metadata_service.py b/surfsense_backend/app/services/google_drive/tool_metadata_service.py new file mode 100644 index 000000000..d7b0a01c3 --- /dev/null +++ b/surfsense_backend/app/services/google_drive/tool_metadata_service.py @@ -0,0 +1,147 @@ +from dataclasses import dataclass + +from sqlalchemy import and_, func +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.future import select + +from app.db import ( + Document, + DocumentType, + SearchSourceConnector, + SearchSourceConnectorType, +) + + +@dataclass +class GoogleDriveAccount: + id: int + name: str + + @classmethod + def from_connector(cls, connector: SearchSourceConnector) -> "GoogleDriveAccount": + return cls(id=connector.id, name=connector.name) + + def to_dict(self) -> dict: + return {"id": self.id, "name": self.name} + + +@dataclass +class GoogleDriveFile: + file_id: str + name: str + mime_type: str + web_view_link: str + connector_id: int + document_id: int + + @classmethod + def from_document(cls, document: Document) -> "GoogleDriveFile": + meta = document.document_metadata or {} + return cls( + file_id=meta.get("google_drive_file_id", ""), + name=meta.get("google_drive_file_name", document.title), + mime_type=meta.get("google_drive_mime_type", ""), + web_view_link=meta.get("web_view_link", ""), + connector_id=document.connector_id, + document_id=document.id, + ) + + def to_dict(self) -> dict: + return { + "file_id": self.file_id, + "name": self.name, + "mime_type": self.mime_type, + "web_view_link": self.web_view_link, + "connector_id": self.connector_id, + "document_id": self.document_id, + } + + +class GoogleDriveToolMetadataService: + def __init__(self, db_session: AsyncSession): + self._db_session = db_session + + async def get_creation_context(self, search_space_id: int, user_id: str) -> dict: + accounts = await self._get_google_drive_accounts(search_space_id, user_id) + + if not accounts: + return { + "accounts": [], + "supported_types": [], + "error": "No Google Drive account connected", + } + + return { + "accounts": [acc.to_dict() for acc in accounts], + "supported_types": ["google_doc", "google_sheet"], + } + + async def get_trash_context( + self, search_space_id: int, user_id: str, file_name: str + ) -> dict: + result = await self._db_session.execute( + select(Document) + .join( + SearchSourceConnector, Document.connector_id == SearchSourceConnector.id + ) + .filter( + and_( + Document.search_space_id == search_space_id, + Document.document_type == DocumentType.GOOGLE_DRIVE_FILE, + func.lower(Document.title) == func.lower(file_name), + SearchSourceConnector.user_id == user_id, + ) + ) + ) + document = result.scalars().first() + + if not document: + return { + "error": ( + f"File '{file_name}' not found in your indexed Google Drive files. " + "This could mean: (1) the file doesn't exist, (2) it hasn't been indexed yet, " + "or (3) the file name is different." + ) + } + + if not document.connector_id: + return {"error": "Document has no associated connector"} + + result = await self._db_session.execute( + select(SearchSourceConnector).filter( + and_( + SearchSourceConnector.id == document.connector_id, + SearchSourceConnector.user_id == user_id, + ) + ) + ) + connector = result.scalars().first() + + if not connector: + return {"error": "Connector not found or access denied"} + + account = GoogleDriveAccount.from_connector(connector) + file = GoogleDriveFile.from_document(document) + + return { + "account": account.to_dict(), + "file": file.to_dict(), + } + + async def _get_google_drive_accounts( + self, search_space_id: int, user_id: str + ) -> list[GoogleDriveAccount]: + result = await self._db_session.execute( + select(SearchSourceConnector) + .filter( + and_( + SearchSourceConnector.search_space_id == search_space_id, + SearchSourceConnector.user_id == user_id, + SearchSourceConnector.connector_type + == SearchSourceConnectorType.GOOGLE_DRIVE_CONNECTOR, + ) + ) + .order_by(SearchSourceConnector.last_indexed_at.desc()) + ) + connectors = result.scalars().all() + return [GoogleDriveAccount.from_connector(c) for c in connectors]