From f98ab2367a21d74c5dad3a0eba39f86f9dd869d4 Mon Sep 17 00:00:00 2001
From: Mike Fortman <michael.fortman@datastax.com>
Date: Wed, 3 Sep 2025 16:20:30 -0500
Subject: [PATCH] add google drive file/folder selector

---
 frontend/src/app/settings/page.tsx            |  73 ++++-
 .../src/components/google-drive-picker.tsx    | 269 ++++++++++++++++++
 src/api/connectors.py                         |  40 +++
 src/main.py                                   |  11 +
 4 files changed, 389 insertions(+), 4 deletions(-)
 create mode 100644 frontend/src/components/google-drive-picker.tsx

diff --git a/frontend/src/app/settings/page.tsx b/frontend/src/app/settings/page.tsx
index c42cbeb8..a1eea2a7 100644
--- a/frontend/src/app/settings/page.tsx
+++ b/frontend/src/app/settings/page.tsx
@@ -12,8 +12,17 @@ import { Loader2, PlugZap, RefreshCw } from "lucide-react"
 import { ProtectedRoute } from "@/components/protected-route"
 import { useTask } from "@/contexts/task-context"
 import { useAuth } from "@/contexts/auth-context"
+import { GoogleDrivePicker } from "@/components/google-drive-picker"
 
 
+interface GoogleDriveFile {
+  id: string
+  name: string
+  mimeType: string
+  webViewLink?: string
+  iconLink?: string
+}
+
 interface Connector {
   id: string
   name: string
@@ -23,6 +32,7 @@ interface Connector {
   type: string
   connectionId?: string
   access_token?: string
+  selectedFiles?: GoogleDriveFile[]
 }
 
 interface SyncResult {
@@ -53,6 +63,8 @@ function KnowledgeSourcesPage() {
   const [syncResults, setSyncResults] = useState<{[key: string]: SyncResult | null}>({})
   const [maxFiles, setMaxFiles] = useState<number>(10)
   const [syncAllFiles, setSyncAllFiles] = useState<boolean>(false)
+  const [selectedFiles, setSelectedFiles] = useState<{[connectorId: string]: GoogleDriveFile[]}>({})
+  const [connectorAccessTokens, setConnectorAccessTokens] = useState<{[connectorId: string]: string}>({})
   
   // Settings state
   // Note: backend internal Langflow URL is not needed on the frontend
@@ -143,6 +155,24 @@ function KnowledgeSourcesPage() {
           const activeConnection = connections.find((conn: Connection) => conn.is_active)
           const isConnected = activeConnection !== undefined
           
+          // For Google Drive, try to get access token for the picker
+          if (connectorType === 'google_drive' && activeConnection) {
+            try {
+              const tokenResponse = await fetch(`/api/connectors/${connectorType}/token?connection_id=${activeConnection.connection_id}`)
+              if (tokenResponse.ok) {
+                const tokenData = await tokenResponse.json()
+                if (tokenData.access_token) {
+                  setConnectorAccessTokens(prev => ({
+                    ...prev,
+                    [connectorType]: tokenData.access_token
+                  }))
+                }
+              }
+            } catch (e) {
+              console.log('Could not fetch access token for Google Drive picker:', e)
+            }
+          }
+          
           setConnectors(prev => prev.map(c => 
             c.type === connectorType 
               ? { 
@@ -208,6 +238,20 @@ function KnowledgeSourcesPage() {
     }
   }
 
+  const handleFileSelection = (connectorId: string, files: GoogleDriveFile[]) => {
+    setSelectedFiles(prev => ({
+      ...prev,
+      [connectorId]: files
+    }))
+    
+    // Update the connector with selected files
+    setConnectors(prev => prev.map(c => 
+      c.id === connectorId 
+        ? { ...c, selectedFiles: files }
+        : c
+    ))
+  }
+
   const handleSync = async (connector: Connector) => {
     if (!connector.connectionId) return
     
@@ -215,15 +259,26 @@ function KnowledgeSourcesPage() {
     setSyncResults(prev => ({ ...prev, [connector.id]: null }))
     
     try {
+      const syncBody: {
+        connection_id: string;
+        max_files?: number;
+        selected_files?: string[];
+      } = {
+        connection_id: connector.connectionId,
+        max_files: syncAllFiles ? 0 : (maxFiles || undefined)
+      }
+      
+      // Add selected files for Google Drive
+      if (connector.type === "google_drive" && selectedFiles[connector.id]?.length > 0) {
+        syncBody.selected_files = selectedFiles[connector.id].map(file => file.id)
+      }
+      
       const response = await fetch(`/api/connectors/${connector.type}/sync`, {
         method: 'POST',
         headers: {
           'Content-Type': 'application/json',
         },
-        body: JSON.stringify({
-          connection_id: connector.connectionId,
-          max_files: syncAllFiles ? 0 : (maxFiles || undefined)
-        }),
+        body: JSON.stringify(syncBody),
       })
       
       const result = await response.json()
@@ -433,6 +488,16 @@ function KnowledgeSourcesPage() {
               <CardContent className="flex-1 flex flex-col justify-end space-y-4">
                 {connector.status === "connected" ? (
                   <div className="space-y-3">
+                    {/* Google Drive file picker */}
+                    {connector.type === "google_drive" && (
+                      <GoogleDrivePicker
+                        onFileSelected={(files) => handleFileSelection(connector.id, files)}
+                        selectedFiles={selectedFiles[connector.id] || []}
+                        isAuthenticated={connector.status === "connected"}
+                        accessToken={connectorAccessTokens[connector.type]}
+                      />
+                    )}
+                    
                     <Button
                       onClick={() => handleSync(connector)}
                       disabled={isSyncing === connector.id}
diff --git a/frontend/src/components/google-drive-picker.tsx b/frontend/src/components/google-drive-picker.tsx
new file mode 100644
index 00000000..50d789a7
--- /dev/null
+++ b/frontend/src/components/google-drive-picker.tsx
@@ -0,0 +1,269 @@
+"use client"
+
+import { useState, useEffect } from "react"
+import { Button } from "@/components/ui/button"
+import { Badge } from "@/components/ui/badge"
+import { FileText, Folder, X } from "lucide-react"
+
+interface GoogleDrivePickerProps {
+  onFileSelected: (files: GoogleDriveFile[]) => void
+  selectedFiles?: GoogleDriveFile[]
+  isAuthenticated: boolean
+  accessToken?: string
+}
+
+interface GoogleDriveFile {
+  id: string
+  name: string
+  mimeType: string
+  webViewLink?: string
+  iconLink?: string
+}
+
+interface GoogleAPI {
+  load: (api: string, options: { callback: () => void; onerror?: () => void }) => void
+}
+
+interface GooglePickerData {
+  action: string
+  docs: GooglePickerDocument[]
+}
+
+interface GooglePickerDocument {
+  [key: string]: string
+}
+
+declare global {
+  interface Window {
+    gapi: GoogleAPI
+    google: {
+      picker: {
+        api: {
+          load: (callback: () => void) => void
+        }
+        PickerBuilder: new () => GooglePickerBuilder
+        ViewId: {
+          DOCS: string
+          FOLDERS: string
+          DOCS_IMAGES_AND_VIDEOS: string
+          DOCUMENTS: string
+          PRESENTATIONS: string
+          SPREADSHEETS: string
+        }
+        Feature: {
+          MULTISELECT_ENABLED: string
+          NAV_HIDDEN: string
+          SIMPLE_UPLOAD_ENABLED: string
+        }
+        Action: {
+          PICKED: string
+          CANCEL: string
+        }
+        Document: {
+          ID: string
+          NAME: string
+          MIME_TYPE: string
+          URL: string
+          ICON_URL: string
+        }
+      }
+    }
+  }
+}
+
+interface GooglePickerBuilder {
+  addView: (view: string) => GooglePickerBuilder
+  setOAuthToken: (token: string) => GooglePickerBuilder
+  setCallback: (callback: (data: GooglePickerData) => void) => GooglePickerBuilder
+  enableFeature: (feature: string) => GooglePickerBuilder
+  setTitle: (title: string) => GooglePickerBuilder
+  build: () => GooglePicker
+}
+
+interface GooglePicker {
+  setVisible: (visible: boolean) => void
+}
+
+export function GoogleDrivePicker({ 
+  onFileSelected, 
+  selectedFiles = [], 
+  isAuthenticated, 
+  accessToken 
+}: GoogleDrivePickerProps) {
+  const [isPickerLoaded, setIsPickerLoaded] = useState(false)
+  const [isPickerOpen, setIsPickerOpen] = useState(false)
+  
+  useEffect(() => {
+    const loadPickerApi = () => {
+      if (typeof window !== 'undefined' && window.gapi) {
+        window.gapi.load('picker', {
+          callback: () => {
+            setIsPickerLoaded(true)
+          },
+          onerror: () => {
+            console.error('Failed to load Google Picker API')
+          }
+        })
+      }
+    }
+
+    // Load Google API script if not already loaded
+    if (typeof window !== 'undefined') {
+      if (!window.gapi) {
+        const script = document.createElement('script')
+        script.src = 'https://apis.google.com/js/api.js'
+        script.async = true
+        script.defer = true
+        script.onload = loadPickerApi
+        script.onerror = () => {
+          console.error('Failed to load Google API script')
+        }
+        document.head.appendChild(script)
+        
+        return () => {
+          if (document.head.contains(script)) {
+            document.head.removeChild(script)
+          }
+        }
+      } else {
+        loadPickerApi()
+      }
+    }
+  }, [])
+
+  const openPicker = () => {
+    if (!isPickerLoaded || !accessToken || !window.google?.picker) {
+      return
+    }
+
+    try {
+      setIsPickerOpen(true)
+
+      const picker = new window.google.picker.PickerBuilder()
+        .addView(window.google.picker.ViewId.DOCS)
+        .addView(window.google.picker.ViewId.FOLDERS)
+        .setOAuthToken(accessToken)
+        .enableFeature(window.google.picker.Feature.MULTISELECT_ENABLED)
+        .setTitle('Select files from Google Drive')
+        .setCallback(pickerCallback)
+        .build()
+
+      picker.setVisible(true)
+    } catch (error) {
+      console.error('Error creating picker:', error)
+      setIsPickerOpen(false)
+    }
+  }
+
+  const pickerCallback = (data: GooglePickerData) => {
+    if (data.action === window.google.picker.Action.PICKED) {
+      const files: GoogleDriveFile[] = data.docs.map((doc: GooglePickerDocument) => ({
+        id: doc[window.google.picker.Document.ID],
+        name: doc[window.google.picker.Document.NAME],
+        mimeType: doc[window.google.picker.Document.MIME_TYPE],
+        webViewLink: doc[window.google.picker.Document.URL],
+        iconLink: doc[window.google.picker.Document.ICON_URL]
+      }))
+      
+      onFileSelected(files)
+    }
+    
+    setIsPickerOpen(false)
+  }
+
+  const removeFile = (fileId: string) => {
+    const updatedFiles = selectedFiles.filter(file => file.id !== fileId)
+    onFileSelected(updatedFiles)
+  }
+
+  const getFileIcon = (mimeType: string) => {
+    if (mimeType.includes('folder')) {
+      return <Folder className="h-4 w-4" />
+    }
+    return <FileText className="h-4 w-4" />
+  }
+
+  const getMimeTypeLabel = (mimeType: string) => {
+    const typeMap: { [key: string]: string } = {
+      'application/vnd.google-apps.document': 'Google Doc',
+      'application/vnd.google-apps.spreadsheet': 'Google Sheet',
+      'application/vnd.google-apps.presentation': 'Google Slides',
+      'application/vnd.google-apps.folder': 'Folder',
+      'application/pdf': 'PDF',
+      'text/plain': 'Text',
+      'application/vnd.openxmlformats-officedocument.wordprocessingml.document': 'Word Doc',
+      'application/vnd.openxmlformats-officedocument.presentationml.presentation': 'PowerPoint'
+    }
+    
+    return typeMap[mimeType] || 'Document'
+  }
+
+  if (!isAuthenticated) {
+    return (
+      <div className="text-sm text-muted-foreground p-4 bg-muted/20 rounded-md">
+        Please connect to Google Drive first to select specific files.
+      </div>
+    )
+  }
+
+  return (
+    <div className="space-y-4">
+      <div className="flex items-center justify-between">
+        <div>
+          <h4 className="text-sm font-medium">File Selection</h4>
+          <p className="text-xs text-muted-foreground">
+            Choose specific files to sync instead of syncing everything
+          </p>
+        </div>
+        <Button
+          onClick={openPicker}
+          disabled={!isPickerLoaded || isPickerOpen || !accessToken}
+          size="sm"
+          variant="outline"
+        >
+          {isPickerOpen ? 'Opening Picker...' : 'Select Files'}
+        </Button>
+      </div>
+
+      {selectedFiles.length > 0 && (
+        <div className="space-y-2">
+          <p className="text-xs text-muted-foreground">
+            Selected files ({selectedFiles.length}):
+          </p>
+          <div className="max-h-32 overflow-y-auto space-y-1">
+            {selectedFiles.map((file) => (
+              <div
+                key={file.id}
+                className="flex items-center justify-between p-2 bg-muted/30 rounded-md text-xs"
+              >
+                <div className="flex items-center gap-2 flex-1 min-w-0">
+                  {getFileIcon(file.mimeType)}
+                  <span className="truncate font-medium">{file.name}</span>
+                  <Badge variant="secondary" className="text-xs px-1 py-0.5 h-auto">
+                    {getMimeTypeLabel(file.mimeType)}
+                  </Badge>
+                </div>
+                <Button
+                  onClick={() => removeFile(file.id)}
+                  size="sm"
+                  variant="ghost"
+                  className="h-6 w-6 p-0"
+                >
+                  <X className="h-3 w-3" />
+                </Button>
+              </div>
+            ))}
+          </div>
+          <Button
+            onClick={() => onFileSelected([])}
+            size="sm"
+            variant="ghost"
+            className="text-xs h-6"
+          >
+            Clear all
+          </Button>
+        </div>
+      )}
+    </div>
+  )
+}
\ No newline at end of file
diff --git a/src/api/connectors.py b/src/api/connectors.py
index 87f21b4b..6b4e4f5b 100644
--- a/src/api/connectors.py
+++ b/src/api/connectors.py
@@ -272,3 +272,43 @@ async def connector_webhook(request: Request, connector_service, session_manager
         return JSONResponse(
             {"error": f"Webhook processing failed: {str(e)}"}, status_code=500
         )
+
+
+async def connector_token(request: Request, connector_service, session_manager):
+    """Get access token for connector API calls (e.g., Google Picker)"""
+    connector_type = request.path_params.get("connector_type")
+    connection_id = request.query_params.get("connection_id")
+    
+    if not connection_id:
+        return JSONResponse({"error": "connection_id is required"}, status_code=400)
+    
+    user = request.state.user
+    
+    try:
+        # Get the connection and verify it belongs to the user
+        connection = await connector_service.connection_manager.get_connection(connection_id)
+        if not connection or connection.user_id != user.user_id:
+            return JSONResponse({"error": "Connection not found"}, status_code=404)
+        
+        # Get the connector instance
+        connector = await connector_service._get_connector(connection_id)
+        if not connector:
+            return JSONResponse({"error": "Connector not available"}, status_code=404)
+        
+        # For Google Drive, get the access token
+        if connector_type == "google_drive" and hasattr(connector, 'oauth'):
+            await connector.oauth.load_credentials()
+            if connector.oauth.creds and connector.oauth.creds.valid:
+                return JSONResponse({
+                    "access_token": connector.oauth.creds.token,
+                    "expires_in": (connector.oauth.creds.expiry.timestamp() - 
+                                 __import__('time').time()) if connector.oauth.creds.expiry else None
+                })
+            else:
+                return JSONResponse({"error": "Invalid or expired credentials"}, status_code=401)
+        
+        return JSONResponse({"error": "Token not available for this connector type"}, status_code=400)
+        
+    except Exception as e:
+        print(f"Error getting connector token: {e}")
+        return JSONResponse({"error": str(e)}, status_code=500)
diff --git a/src/main.py b/src/main.py
index 611d9885..d1c34546 100644
--- a/src/main.py
+++ b/src/main.py
@@ -586,6 +586,17 @@ async def create_app():
             ),
             methods=["GET"],
         ),
+        Route(
+            "/connectors/{connector_type}/token",
+            require_auth(services["session_manager"])(
+                partial(
+                    connectors.connector_token,
+                    connector_service=services["connector_service"],
+                    session_manager=services["session_manager"],
+                )
+            ),
+            methods=["GET"],
+        ),
         Route(
             "/connectors/{connector_type}/webhook",
             partial(