google-workspace-mcp 1.0.5__py3-none-any.whl → 1.2.0__py3-none-any.whl

google_workspace_mcp/services/drive.py

@@ -27,8 +27,137 @@ class DriveService(BaseGoogleService):
         """Initialize the Drive service."""
         super().__init__("drive", "v3")
 
+    def _escape_drive_query(self, query: str) -> str:
+        """
+        Smart query processing for Drive API queries.
+
+        Automatically converts simple text searches to proper Drive API format:
+        - Simple text like "sprint planning" becomes fullText contains 'sprint planning'
+        - Already formatted queries like "name contains 'test'" are passed through
+        - Handles apostrophe escaping automatically
+        - Validates queries to prevent invalid syntax
+
+        Args:
+            query: Query string (can be simple text or Drive API format)
+
+        Returns:
+            Properly formatted query string for Drive API
+
+        Raises:
+            ValueError: If query contains invalid syntax like unsupported parentheses
+        """
+        if not query or not query.strip():
+            return ""
+
+        query = query.strip()
+
+        # Remove surrounding double quotes if present
+        if query.startswith('"') and query.endswith('"'):
+            query = query[1:-1]
+
+        # Check for unsupported syntax - parentheses are not supported in files.list queries
+        if "(" in query or ")" in query:
+            raise ValueError(
+                "Parentheses are not supported in Drive file search queries. "
+                "Use separate search calls or restructure your query. "
+                "For complex searches, try: fullText contains 'term1' OR fullText contains 'term2'"
+            )
+
+        # Check if this is already a structured Drive API query
+        # Look for Drive API operators and syntax
+        drive_api_indicators = [
+            " contains ",
+            " = ",
+            " != ",
+            " > ",
+            " < ",
+            " >= ",
+            " <= ",
+            " in ",
+            " and ",
+            " or ",
+            " not ",
+            "mimeType",
+            "parents",
+            "trashed",
+            "sharedWithMe",
+            "starred",
+            "modifiedTime",
+            "createdTime",
+            "name",
+            "fullText",
+        ]
+
+        is_structured_query = any(
+            indicator in query.lower() for indicator in drive_api_indicators
+        )
+
+        if is_structured_query:
+            # Validate that it's not a mixed query (text + operators + API syntax)
+            # Mixed queries like "ArcLio (sprint OR planning) modifiedTime > '2024-01-01'" are invalid
+            words = query.split()
+            has_unquoted_text = False
+
+            for word in words:
+                # Skip quoted strings, operators, and API field names
+                if (
+                    word.startswith("'")
+                    or word.endswith("'")
+                    or word.lower() in ["and", "or", "not", "contains", "in"]
+                    or any(
+                        field in word.lower()
+                        for field in [
+                            "mimetype",
+                            "name",
+                            "fulltext",
+                            "modifiedtime",
+                            "createdtime",
+                            "trashed",
+                            "starred",
+                            "sharedwithme",
+                            "parents",
+                        ]
+                    )
+                    or word in ["=", "!=", ">", "<", ">=", "<="]
+                    or word.startswith("'")
+                    and word.endswith("'")
+                ):
+                    continue
+
+                # Check if this looks like unquoted text mixed with operators
+                if not word.replace("-", "").replace("_", "").isalnum():
+                    continue
+
+                # If we find unquoted alphanumeric text in a structured query, it might be mixed
+                if (
+                    any(op in query.lower() for op in [" and ", " or "])
+                    and len(word) > 2
+                ):
+                    has_unquoted_text = True
+                    break
+
+            if has_unquoted_text:
+                raise ValueError(
+                    "Mixed query syntax detected. Use either: "
+                    "1) Simple text: 'sprint planning meeting' "
+                    "2) Drive API syntax: fullText contains 'sprint' OR fullText contains 'planning' "
+                    "3) Structured queries: modifiedTime > '2024-01-01' AND name contains 'sprint'"
+                )
+
+            # This looks like a valid structured Drive API query, just escape apostrophes
+            return query.replace("'", "\\'")
+        # This is a simple text search, convert to fullText search
+        # Escape apostrophes in the search term
+        escaped_text = query.replace("'", "\\'")
+        # Wrap in fullText contains for better search results
+        return f"fullText contains '{escaped_text}'"
+
     def search_files(
-        self, query: str, page_size: int = 10, shared_drive_id: str | None = None
+        self,
+        query: str,
+        page_size: int = 10,
+        shared_drive_id: str | None = None,
+        include_shared_drives: bool = True,
     ) -> list[dict[str, Any]]:
         """
         Search for files in Google Drive.
@@ -37,36 +166,42 @@ class DriveService(BaseGoogleService):
             query: Search query string
             page_size: Maximum number of files to return (1-1000)
             shared_drive_id: Optional shared drive ID to search within a specific shared drive
+            include_shared_drives: Whether to include shared drives in search (default True)
 
         Returns:
             List of file metadata dictionaries (id, name, mimeType, etc.) or an error dictionary
         """
         try:
             logger.info(
-                f"Searching files with query: '{query}', page_size: {page_size}, shared_drive_id: {shared_drive_id}"
+                f"Searching files with query: '{query}', page_size: {page_size}, "
+                f"shared_drive_id: {shared_drive_id}, include_shared_drives: {include_shared_drives}"
             )
 
             # Validate and constrain page_size
             page_size = max(1, min(page_size, 1000))
 
-            # Build list parameters with shared drive support
+            # Properly escape the query for Drive API
+            escaped_query = self._escape_drive_query(query)
+
+            # Build list parameters with comprehensive shared drive support
             list_params = {
-                "q": query,  # Use the query directly without modification
+                "q": escaped_query,
                 "pageSize": page_size,
-                "fields": "files(id, name, mimeType, modifiedTime, size, webViewLink, iconLink)",
+                "fields": "files(id, name, mimeType, modifiedTime, size, webViewLink, iconLink, parents)",
                 "supportsAllDrives": True,
                 "includeItemsFromAllDrives": True,
             }
 
             if shared_drive_id:
+                # Search within a specific shared drive
                 list_params["driveId"] = shared_drive_id
-                list_params["corpora"] = (
-                    "drive"  # Search within the specified shared drive
-                )
+                list_params["corpora"] = "drive"
+            elif include_shared_drives:
+                # Search across all drives (user's files + shared drives + shared folders)
+                list_params["corpora"] = "allDrives"
             else:
-                list_params["corpora"] = (
-                    "user"  # Default to user's files if no specific shared drive ID
-                )
+                # Search only user's personal files
+                list_params["corpora"] = "user"
 
             results = self.service.files().list(**list_params).execute()
             files = results.get("files", [])
@@ -91,7 +226,7 @@ class DriveService(BaseGoogleService):
             # Get file metadata
             file_metadata = (
                 self.service.files()
-                .get(fileId=file_id, fields="mimeType, name")
+                .get(fileId=file_id, fields="mimeType, name", supportsAllDrives=True)
                 .execute()
             )
 
@@ -246,21 +381,27 @@ class DriveService(BaseGoogleService):
                 try:
                     content = content_bytes.decode("utf-8")
                     return {
-                        "mimeType": export_mime_type,
+                        "file_id": file_id,
+                        "name": file_name,
+                        "mime_type": export_mime_type,
                         "content": content,
                         "encoding": "utf-8",
                     }
                 except UnicodeDecodeError:
                     content = base64.b64encode(content_bytes).decode("utf-8")
                     return {
-                        "mimeType": export_mime_type,
+                        "file_id": file_id,
+                        "name": file_name,
+                        "mime_type": export_mime_type,
                         "content": content,
                         "encoding": "base64",
                     }
             else:
                 content = base64.b64encode(content_bytes).decode("utf-8")
                 return {
-                    "mimeType": export_mime_type,
+                    "file_id": file_id,
+                    "name": file_name,
+                    "mime_type": export_mime_type,
                     "content": content,
                     "encoding": "base64",
                 }
@@ -281,21 +422,35 @@ class DriveService(BaseGoogleService):
         if mime_type.startswith("text/") or mime_type == "application/json":
             try:
                 content = content_bytes.decode("utf-8")
-                return {"mimeType": mime_type, "content": content, "encoding": "utf-8"}
+                return {
+                    "file_id": file_id,
+                    "name": file_name,
+                    "mime_type": mime_type,
+                    "content": content,
+                    "encoding": "utf-8",
+                }
             except UnicodeDecodeError:
                 logger.warning(
                     f"UTF-8 decoding failed for file {file_id} ('{file_name}', {mime_type}). Using base64."
                 )
                 content = base64.b64encode(content_bytes).decode("utf-8")
                 return {
-                    "mimeType": mime_type,
+                    "file_id": file_id,
+                    "name": file_name,
+                    "mime_type": mime_type,
                     "content": content,
                     "encoding": "base64",
                 }
         else:
             # Binary file
             content = base64.b64encode(content_bytes).decode("utf-8")
-            return {"mimeType": mime_type, "content": content, "encoding": "base64"}
+            return {
+                "file_id": file_id,
+                "name": file_name,
+                "mime_type": mime_type,
+                "content": content,
+                "encoding": "base64",
+            }
 
     def _download_content(self, request) -> bytes:
         """Download content from a request."""
@@ -448,3 +603,98 @@ class DriveService(BaseGoogleService):
                 "message": str(e),
                 "operation": "list_shared_drives",
             }
+
+    def share_file_with_domain(
+        self, file_id: str, domain: str, role: str = "reader"
+    ) -> dict[str, Any]:
+        """
+        Shares a file with an entire domain.
+
+        Args:
+            file_id: The ID of the file to share.
+            domain: The domain to share the file with (e.g., 'example.com').
+            role: The permission role ('reader', 'commenter', 'writer'). Defaults to 'reader'.
+
+        Returns:
+            A dictionary containing the permission details or an error dictionary.
+        """
+        try:
+            if not file_id or not domain:
+                raise ValueError("File ID and domain are required.")
+
+            logger.info(
+                f"Sharing file {file_id} with domain '{domain}' as role '{role}'"
+            )
+
+            permission = {"type": "domain", "role": role, "domain": domain}
+
+            # Create the permission
+            permission_result = (
+                self.service.permissions()
+                .create(
+                    fileId=file_id,
+                    body=permission,
+                    sendNotificationEmail=False,  # Avoid spamming the domain
+                    supportsAllDrives=True,
+                )
+                .execute()
+            )
+
+            logger.info(
+                f"Successfully created domain permission ID: {permission_result.get('id')}"
+            )
+            return {
+                "success": True,
+                "file_id": file_id,
+                "permission_id": permission_result.get("id"),
+                "domain": domain,
+                "role": role,
+            }
+
+        except HttpError as error:
+            # Check for a 403 error related to sharing policies
+            if error.resp.status == 403:
+                error_content = error.content.decode("utf-8")
+                if (
+                    "cannotShareTeamDriveItem" in error_content
+                    or "sharingRateLimitExceeded" in error_content
+                ):
+                    logger.error(
+                        f"Domain sharing policy prevents sharing file {file_id}: {error_content}"
+                    )
+                    # Return a more specific error message
+                    return self.handle_api_error(
+                        "share_file_with_domain_policy_error", error
+                    )
+
+            return self.handle_api_error("share_file_with_domain", error)
+        except Exception as e:
+            return self.handle_api_error("share_file_with_domain", e)
+
+    def _get_or_create_data_folder(self) -> str:
+        """
+        Finds the dedicated folder for storing chart data, creating it if it doesn't exist.
+        The result is cached to avoid repeated API calls within the same session.
+
+        Returns:
+            The ID of the data folder.
+        """
+        folder_name = "[MCP] Generated Chart Data"
+        query = f"name = '{folder_name}' and mimeType = 'application/vnd.google-apps.folder' and trashed = false"
+
+        logger.info(f"Searching for data folder: '{folder_name}'")
+        search_result = self.search_files(query=query, page_size=1)
+
+        if search_result and len(search_result) > 0:
+            folder_id = search_result[0]["id"]
+            logger.info(f"Found existing data folder with ID: {folder_id}")
+            return folder_id
+        logger.info("Data folder not found. Creating a new one.")
+        create_result = self.create_folder(folder_name=folder_name)
+        if create_result and not create_result.get("error"):
+            folder_id = create_result["id"]
+            logger.info(f"Successfully created data folder with ID: {folder_id}")
+            return folder_id
+        raise RuntimeError(
+            "Failed to create the necessary data folder in Google Drive."
+        )