diff --git a/README.md b/README.md index 23f9ccc..cdf04b0 100644 --- a/README.md +++ b/README.md @@ -555,6 +555,29 @@ gro files tree --depth 3 gro drive tree --files # Include files, not just folders ``` +#### Shared Drives + +gro supports Google Shared Drives (formerly Team Drives). By default, search includes files from all drives you have access to. + +```bash +# List available shared drives +gro drive drives +gro drive drives --json + +# Search all drives (default) +gro drive search "quarterly report" + +# Search only your personal drive +gro drive search "quarterly report" --my-drive + +# Search a specific shared drive by name +gro drive search "budget" --drive "Finance Team" +gro drive list --drive "Engineering" +gro drive tree --drive "Marketing" +``` + +The `--my-drive` and `--drive` flags are mutually exclusive. Shared drive names are cached locally for fast lookups. Run `gro drive drives` to refresh the cache. + ### gro drive list List files in Google Drive root or a specific folder. @@ -567,12 +590,14 @@ Aliases: gro files list Flags: -m, --max int Maximum number of files (default 25) -t, --type string Filter by type (document, spreadsheet, presentation, folder, pdf, image, video, audio) + --my-drive List from My Drive only + --drive string List from specific shared drive (name or ID) -j, --json Output as JSON ``` ### gro drive search -Search for files in Google Drive. +Search for files in Google Drive. By default, searches all drives you have access to. ``` Usage: gro drive search [query] [flags] @@ -586,6 +611,8 @@ Flags: --modified-after string Modified after date (YYYY-MM-DD) --modified-before string Modified before date (YYYY-MM-DD) --in-folder string Search within folder ID + --my-drive Search only My Drive + --drive string Search specific shared drive (name or ID) -m, --max int Maximum results (default 25) -j, --json Output as JSON ``` @@ -636,9 +663,25 @@ Aliases: gro files tree Flags: -d, --depth int Maximum depth to traverse (default 2) --files Include files in addition to folders + --my-drive Show My Drive only (default) + --drive string Show tree from specific shared drive -j, --json Output as JSON ``` +### gro drive drives + +List all shared drives accessible to you. + +``` +Usage: gro drive drives [flags] + +Aliases: gro files drives + +Flags: + --refresh Force refresh from API (ignore cache) + -j, --json Output as JSON +``` + ## Search Query Reference gro supports all Gmail search operators: diff --git a/internal/cmd/drive/drive.go b/internal/cmd/drive/drive.go index f01240c..6d08766 100644 --- a/internal/cmd/drive/drive.go +++ b/internal/cmd/drive/drive.go @@ -18,12 +18,20 @@ This command group provides Google Drive functionality: - get: Get detailed metadata for a file - download: Download files or export Google Docs - tree: Display folder structure +- drives: List accessible shared drives + +Shared Drive Support: + By default, search includes files from all drives (My Drive + shared drives). + Use --my-drive to limit to personal drive, or --drive to target a + specific shared drive. Examples: gro drive list gro drive search "quarterly report" + gro drive search "budget" --drive "Finance Team" gro drive get - gro drive download --format pdf`, + gro drive download --format pdf + gro drive drives`, } cmd.AddCommand(newListCommand()) @@ -31,6 +39,7 @@ Examples: cmd.AddCommand(newGetCommand()) cmd.AddCommand(newDownloadCommand()) cmd.AddCommand(newTreeCommand()) + cmd.AddCommand(newDrivesCommand()) return cmd } diff --git a/internal/cmd/drive/drives.go b/internal/cmd/drive/drives.go new file mode 100644 index 0000000..17be98f --- /dev/null +++ b/internal/cmd/drive/drives.go @@ -0,0 +1,179 @@ +package drive + +import ( + "fmt" + "os" + "strings" + "text/tabwriter" + + "github.com/spf13/cobra" + + "github.com/open-cli-collective/google-readonly/internal/cache" + "github.com/open-cli-collective/google-readonly/internal/config" + "github.com/open-cli-collective/google-readonly/internal/drive" +) + +func newDrivesCommand() *cobra.Command { + var ( + jsonOutput bool + refresh bool + ) + + cmd := &cobra.Command{ + Use: "drives", + Short: "List shared drives", + Long: `List all Google Shared Drives accessible to you. + +Results are cached locally for fast lookups. Use --refresh to force a refresh. + +Examples: + gro drive drives # List shared drives (uses cache) + gro drive drives --refresh # Force refresh from API + gro drive drives --json # Output as JSON`, + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, args []string) error { + client, err := newDriveClient() + if err != nil { + return fmt.Errorf("failed to create Drive client: %w", err) + } + + // Initialize cache + ttl := config.GetCacheTTLHours() + c, err := cache.New(ttl) + if err != nil { + return fmt.Errorf("failed to initialize cache: %w", err) + } + + var drives []*drive.SharedDrive + + // Try cache first unless refresh requested + if !refresh { + cached, err := c.GetDrives() + if err != nil { + return fmt.Errorf("failed to read cache: %w", err) + } + if cached != nil { + // Convert from cache type to drive type + drives = make([]*drive.SharedDrive, len(cached)) + for i, d := range cached { + drives[i] = &drive.SharedDrive{ + ID: d.ID, + Name: d.Name, + } + } + } + } + + // Fetch from API if no cache hit + if drives == nil { + drives, err = client.ListSharedDrives(100) + if err != nil { + return fmt.Errorf("failed to list shared drives: %w", err) + } + + // Update cache + cached := make([]*cache.CachedDrive, len(drives)) + for i, d := range drives { + cached[i] = &cache.CachedDrive{ + ID: d.ID, + Name: d.Name, + } + } + if err := c.SetDrives(cached); err != nil { + // Non-fatal: warn but continue + fmt.Fprintf(os.Stderr, "Warning: failed to update cache: %v\n", err) + } + } + + if len(drives) == 0 { + fmt.Println("No shared drives found.") + return nil + } + + if jsonOutput { + return printJSON(drives) + } + + printSharedDrives(drives) + return nil + }, + } + + cmd.Flags().BoolVarP(&jsonOutput, "json", "j", false, "Output results as JSON") + cmd.Flags().BoolVar(&refresh, "refresh", false, "Force refresh from API (ignore cache)") + + return cmd +} + +// printSharedDrives prints shared drives in a formatted table +func printSharedDrives(drives []*drive.SharedDrive) { + w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0) + _, _ = fmt.Fprintln(w, "ID\tNAME") + + for _, d := range drives { + _, _ = fmt.Fprintf(w, "%s\t%s\n", d.ID, d.Name) + } + + _ = w.Flush() +} + +// resolveDriveScope converts command flags to a DriveScope, resolving drive names via cache +func resolveDriveScope(client drive.DriveClientInterface, myDrive bool, driveFlag string) (drive.DriveScope, error) { + // --my-drive flag + if myDrive { + return drive.DriveScope{MyDriveOnly: true}, nil + } + + // No --drive flag: search all drives + if driveFlag == "" { + return drive.DriveScope{AllDrives: true}, nil + } + + // --drive flag provided: resolve name or use as ID + // If it looks like a drive ID (starts with 0A), use directly + if looksLikeDriveID(driveFlag) { + return drive.DriveScope{DriveID: driveFlag}, nil + } + + // Try to resolve name via cache + ttl := config.GetCacheTTLHours() + c, err := cache.New(ttl) + if err != nil { + return drive.DriveScope{}, fmt.Errorf("failed to initialize cache: %w", err) + } + + cached, _ := c.GetDrives() + if cached == nil { + // Cache miss - fetch from API + drives, err := client.ListSharedDrives(100) + if err != nil { + return drive.DriveScope{}, fmt.Errorf("failed to list shared drives: %w", err) + } + + // Update cache + cached = make([]*cache.CachedDrive, len(drives)) + for i, d := range drives { + cached[i] = &cache.CachedDrive{ + ID: d.ID, + Name: d.Name, + } + } + _ = c.SetDrives(cached) // Ignore cache write errors + } + + // Find by name (case-insensitive) + nameLower := strings.ToLower(driveFlag) + for _, d := range cached { + if strings.ToLower(d.Name) == nameLower { + return drive.DriveScope{DriveID: d.ID}, nil + } + } + + return drive.DriveScope{}, fmt.Errorf("shared drive not found: %s", driveFlag) +} + +// looksLikeDriveID returns true if the string appears to be a Drive ID +// Shared drive IDs typically start with "0A" +func looksLikeDriveID(s string) bool { + return len(s) > 10 && strings.HasPrefix(s, "0A") +} diff --git a/internal/cmd/drive/drives_test.go b/internal/cmd/drive/drives_test.go new file mode 100644 index 0000000..6a2d2bc --- /dev/null +++ b/internal/cmd/drive/drives_test.go @@ -0,0 +1,221 @@ +package drive + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/open-cli-collective/google-readonly/internal/drive" + "github.com/open-cli-collective/google-readonly/internal/testutil" +) + +func TestDrivesCommand(t *testing.T) { + cmd := newDrivesCommand() + + t.Run("has correct use", func(t *testing.T) { + assert.Equal(t, "drives", cmd.Use) + }) + + t.Run("requires no arguments", func(t *testing.T) { + err := cmd.Args(cmd, []string{}) + assert.NoError(t, err) + + err = cmd.Args(cmd, []string{"extra"}) + assert.Error(t, err) + }) + + t.Run("has json flag", func(t *testing.T) { + flag := cmd.Flags().Lookup("json") + assert.NotNil(t, flag) + assert.Equal(t, "j", flag.Shorthand) + assert.Equal(t, "false", flag.DefValue) + }) + + t.Run("has refresh flag", func(t *testing.T) { + flag := cmd.Flags().Lookup("refresh") + assert.NotNil(t, flag) + assert.Equal(t, "false", flag.DefValue) + }) + + t.Run("has short description", func(t *testing.T) { + assert.Contains(t, cmd.Short, "shared drives") + }) + + t.Run("has long description", func(t *testing.T) { + assert.Contains(t, cmd.Long, "Shared Drives") + assert.Contains(t, cmd.Long, "cache") + }) +} + +func TestLooksLikeDriveID(t *testing.T) { + tests := []struct { + name string + input string + expected bool + }{ + { + name: "valid shared drive ID", + input: "0ALengineering123456", + expected: true, + }, + { + name: "another valid ID", + input: "0ALsomeDriveIdHere789", + expected: true, + }, + { + name: "short string is not ID", + input: "0AL", + expected: false, + }, + { + name: "regular name is not ID", + input: "Engineering", + expected: false, + }, + { + name: "name with spaces", + input: "Finance Team", + expected: false, + }, + { + name: "starts with 0A but too short", + input: "0Ashort", + expected: false, + }, + { + name: "empty string", + input: "", + expected: false, + }, + { + name: "doesn't start with 0A", + input: "1ALengineering123456", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := looksLikeDriveID(tt.input) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestResolveDriveScope(t *testing.T) { + t.Run("returns MyDriveOnly when myDrive flag is true", func(t *testing.T) { + mock := &testutil.MockDriveClient{} + + scope, err := resolveDriveScope(mock, true, "") + + assert.NoError(t, err) + assert.True(t, scope.MyDriveOnly) + assert.False(t, scope.AllDrives) + assert.Empty(t, scope.DriveID) + }) + + t.Run("returns AllDrives when no flags provided", func(t *testing.T) { + mock := &testutil.MockDriveClient{} + + scope, err := resolveDriveScope(mock, false, "") + + assert.NoError(t, err) + assert.True(t, scope.AllDrives) + assert.False(t, scope.MyDriveOnly) + assert.Empty(t, scope.DriveID) + }) + + t.Run("returns DriveID directly when input looks like ID", func(t *testing.T) { + mock := &testutil.MockDriveClient{} + + scope, err := resolveDriveScope(mock, false, "0ALengineering123456") + + assert.NoError(t, err) + assert.Equal(t, "0ALengineering123456", scope.DriveID) + assert.False(t, scope.AllDrives) + assert.False(t, scope.MyDriveOnly) + }) + + t.Run("resolves drive name to ID via API", func(t *testing.T) { + mock := &testutil.MockDriveClient{ + ListSharedDrivesFunc: func(pageSize int64) ([]*drive.SharedDrive, error) { + return []*drive.SharedDrive{ + {ID: "0ALeng123", Name: "Engineering"}, + {ID: "0ALfin456", Name: "Finance"}, + }, nil + }, + } + + scope, err := resolveDriveScope(mock, false, "Engineering") + + assert.NoError(t, err) + assert.Equal(t, "0ALeng123", scope.DriveID) + }) + + t.Run("resolves drive name case-insensitively", func(t *testing.T) { + mock := &testutil.MockDriveClient{ + ListSharedDrivesFunc: func(pageSize int64) ([]*drive.SharedDrive, error) { + return []*drive.SharedDrive{ + {ID: "0ALeng123", Name: "Engineering"}, + }, nil + }, + } + + scope, err := resolveDriveScope(mock, false, "ENGINEERING") + + assert.NoError(t, err) + assert.Equal(t, "0ALeng123", scope.DriveID) + }) + + t.Run("returns error when drive name not found", func(t *testing.T) { + mock := &testutil.MockDriveClient{ + ListSharedDrivesFunc: func(pageSize int64) ([]*drive.SharedDrive, error) { + return []*drive.SharedDrive{ + {ID: "0ALeng123", Name: "Engineering"}, + }, nil + }, + } + + _, err := resolveDriveScope(mock, false, "NonExistent") + + assert.Error(t, err) + assert.Contains(t, err.Error(), "shared drive not found") + }) +} + +func TestSearchCommand_MutualExclusivity(t *testing.T) { + t.Run("errors when both my-drive and drive flags set", func(t *testing.T) { + cmd := newSearchCommand() + cmd.SetArgs([]string{"query", "--my-drive", "--drive", "Engineering"}) + + err := cmd.Execute() + + assert.Error(t, err) + assert.Contains(t, err.Error(), "mutually exclusive") + }) +} + +func TestListCommand_MutualExclusivity(t *testing.T) { + t.Run("errors when both my-drive and drive flags set", func(t *testing.T) { + cmd := newListCommand() + cmd.SetArgs([]string{"--my-drive", "--drive", "Engineering"}) + + err := cmd.Execute() + + assert.Error(t, err) + assert.Contains(t, err.Error(), "mutually exclusive") + }) +} + +func TestTreeCommand_MutualExclusivity(t *testing.T) { + t.Run("errors when both my-drive and drive flags set", func(t *testing.T) { + cmd := newTreeCommand() + cmd.SetArgs([]string{"--my-drive", "--drive", "Engineering"}) + + err := cmd.Execute() + + assert.Error(t, err) + assert.Contains(t, err.Error(), "mutually exclusive") + }) +} diff --git a/internal/cmd/drive/list.go b/internal/cmd/drive/list.go index 1558385..554f43e 100644 --- a/internal/cmd/drive/list.go +++ b/internal/cmd/drive/list.go @@ -17,6 +17,8 @@ func newListCommand() *cobra.Command { maxResults int64 fileType string jsonOutput bool + myDrive bool + driveFlag string ) cmd := &cobra.Command{ @@ -24,16 +26,25 @@ func newListCommand() *cobra.Command { Short: "List files in Drive", Long: `List files in Google Drive root or a specific folder. +By default, lists files in My Drive root. Use --drive to list files in a +specific shared drive's root. + Examples: - gro drive list # List files in root - gro drive list # List files in specific folder - gro drive list --type document # Filter by file type - gro drive list --max 50 # Limit results - gro drive list --json # Output as JSON + gro drive list # List files in My Drive root + gro drive list # List files in specific folder + gro drive list --drive "Engineering" # List files in shared drive root + gro drive list --type document # Filter by file type + gro drive list --max 50 # Limit results + gro drive list --json # Output as JSON File types: document, spreadsheet, presentation, folder, pdf, image, video, audio`, Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { + // Validate mutually exclusive flags + if myDrive && driveFlag != "" { + return fmt.Errorf("--my-drive and --drive are mutually exclusive") + } + client, err := newDriveClient() if err != nil { return fmt.Errorf("failed to create Drive client: %w", err) @@ -44,12 +55,18 @@ File types: document, spreadsheet, presentation, folder, pdf, image, video, audi folderID = args[0] } - query, err := buildListQuery(folderID, fileType) + // Resolve drive scope for listing + scope, err := resolveDriveScopeForList(client, myDrive, driveFlag, folderID) + if err != nil { + return fmt.Errorf("failed to resolve drive scope: %w", err) + } + + query, err := buildListQueryWithScope(folderID, fileType, scope) if err != nil { return fmt.Errorf("failed to build query: %w", err) } - files, err := client.ListFiles(query, maxResults) + files, err := client.ListFilesWithScope(query, maxResults, scope) if err != nil { return fmt.Errorf("failed to list files: %w", err) } @@ -71,6 +88,8 @@ File types: document, spreadsheet, presentation, folder, pdf, image, video, audi cmd.Flags().Int64VarP(&maxResults, "max", "m", 25, "Maximum number of results to return") cmd.Flags().StringVarP(&fileType, "type", "t", "", "Filter by file type") cmd.Flags().BoolVarP(&jsonOutput, "json", "j", false, "Output results as JSON") + cmd.Flags().BoolVar(&myDrive, "my-drive", false, "Limit to My Drive only") + cmd.Flags().StringVar(&driveFlag, "drive", "", "List files in specific shared drive (name or ID)") return cmd } @@ -96,6 +115,42 @@ func buildListQuery(folderID, fileType string) (string, error) { return strings.Join(parts, " and "), nil } +// buildListQueryWithScope constructs a Drive API query string with scope awareness +func buildListQueryWithScope(folderID, fileType string, scope drive.DriveScope) (string, error) { + parts := []string{"trashed = false"} + + // For shared drives, if no folder specified, we don't add 'root' in parents + // because the root is the drive itself + if folderID != "" { + parts = append(parts, fmt.Sprintf("'%s' in parents", folderID)) + } else if scope.DriveID == "" { + // Only add 'root' for My Drive listings + parts = append(parts, "'root' in parents") + } + + if fileType != "" { + filter, err := getMimeTypeFilter(fileType) + if err != nil { + return "", err + } + parts = append(parts, filter) + } + + return strings.Join(parts, " and "), nil +} + +// resolveDriveScopeForList resolves the scope for list operations +// List has slightly different behavior - defaults to My Drive root if no flags +func resolveDriveScopeForList(client drive.DriveClientInterface, myDrive bool, driveFlag, folderID string) (drive.DriveScope, error) { + // If a folder ID is provided, we need to support all drives to access it + if folderID != "" && !myDrive && driveFlag == "" { + return drive.DriveScope{AllDrives: true}, nil + } + + // Otherwise use the standard resolution + return resolveDriveScope(client, myDrive, driveFlag) +} + // getMimeTypeFilter returns the Drive API query filter for a file type func getMimeTypeFilter(fileType string) (string, error) { switch strings.ToLower(fileType) { diff --git a/internal/cmd/drive/search.go b/internal/cmd/drive/search.go index 8adb217..625a64c 100644 --- a/internal/cmd/drive/search.go +++ b/internal/cmd/drive/search.go @@ -17,6 +17,8 @@ func newSearchCommand() *cobra.Command { modBefore string inFolder string jsonOutput bool + myDrive bool + driveFlag string ) cmd := &cobra.Command{ @@ -24,8 +26,14 @@ func newSearchCommand() *cobra.Command { Short: "Search for files", Long: `Search for files in Google Drive by content, name, type, owner, or date. +By default, searches all drives (My Drive + shared drives you have access to). +Use --my-drive to limit to your personal drive, or --drive to search a specific +shared drive. + Examples: - gro drive search "quarterly report" # Full-text search + gro drive search "quarterly report" # Full-text search (all drives) + gro drive search "quarterly report" --my-drive # Search My Drive only + gro drive search "budget" --drive "Finance" # Search specific shared drive gro drive search --name "budget" # Search by filename only gro drive search --type spreadsheet # Filter by type gro drive search --owner me # Files you own @@ -37,6 +45,11 @@ Examples: File types: document, spreadsheet, presentation, folder, pdf, image, video, audio`, Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { + // Validate mutually exclusive flags + if myDrive && driveFlag != "" { + return fmt.Errorf("--my-drive and --drive are mutually exclusive") + } + client, err := newDriveClient() if err != nil { return fmt.Errorf("failed to create Drive client: %w", err) @@ -52,7 +65,13 @@ File types: document, spreadsheet, presentation, folder, pdf, image, video, audi return fmt.Errorf("failed to build search query: %w", err) } - files, err := client.ListFiles(searchQuery, maxResults) + // Resolve drive scope + scope, err := resolveDriveScope(client, myDrive, driveFlag) + if err != nil { + return fmt.Errorf("failed to resolve drive scope: %w", err) + } + + files, err := client.ListFilesWithScope(searchQuery, maxResults, scope) if err != nil { return fmt.Errorf("failed to search files: %w", err) } @@ -88,6 +107,8 @@ File types: document, spreadsheet, presentation, folder, pdf, image, video, audi cmd.Flags().StringVar(&modBefore, "modified-before", "", "Modified before date (YYYY-MM-DD)") cmd.Flags().StringVar(&inFolder, "in-folder", "", "Search within specific folder") cmd.Flags().BoolVarP(&jsonOutput, "json", "j", false, "Output results as JSON") + cmd.Flags().BoolVar(&myDrive, "my-drive", false, "Limit search to My Drive only") + cmd.Flags().StringVar(&driveFlag, "drive", "", "Search in specific shared drive (name or ID)") return cmd } diff --git a/internal/cmd/drive/tree.go b/internal/cmd/drive/tree.go index b7a9564..c92fc1e 100644 --- a/internal/cmd/drive/tree.go +++ b/internal/cmd/drive/tree.go @@ -22,6 +22,8 @@ func newTreeCommand() *cobra.Command { depth int files bool jsonOutput bool + myDrive bool + driveFlag string ) cmd := &cobra.Command{ @@ -29,26 +31,46 @@ func newTreeCommand() *cobra.Command { Short: "Display folder structure", Long: `Display the folder structure of Google Drive in a tree format. +By default, shows My Drive structure. Use --drive to show a shared drive's +folder structure. + Examples: - gro drive tree # Show folder tree from root - gro drive tree # Show tree from specific folder - gro drive tree --depth 3 # Limit depth - gro drive tree --files # Include files, not just folders - gro drive tree --json # Output as JSON`, + gro drive tree # Show folder tree from My Drive root + gro drive tree # Show tree from specific folder + gro drive tree --drive "Engineering" # Show tree from shared drive root + gro drive tree --depth 3 # Limit depth + gro drive tree --files # Include files, not just folders + gro drive tree --json # Output as JSON`, Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { + // Validate mutually exclusive flags + if myDrive && driveFlag != "" { + return fmt.Errorf("--my-drive and --drive are mutually exclusive") + } + client, err := newDriveClient() if err != nil { return fmt.Errorf("failed to create Drive client: %w", err) } folderID := "root" + rootName := "My Drive" + if len(args) > 0 { folderID = args[0] + rootName = "" // Will be fetched from folder info + } else if driveFlag != "" { + // Resolve shared drive + scope, err := resolveDriveScope(client, false, driveFlag) + if err != nil { + return fmt.Errorf("failed to resolve drive: %w", err) + } + folderID = scope.DriveID + rootName = driveFlag // Use the provided name } // Build the tree - tree, err := buildTree(client, folderID, depth, files) + tree, err := buildTreeWithScope(client, folderID, rootName, depth, files) if err != nil { return fmt.Errorf("failed to build folder tree: %w", err) } @@ -65,12 +87,19 @@ Examples: cmd.Flags().IntVarP(&depth, "depth", "d", 2, "Maximum depth to traverse") cmd.Flags().BoolVar(&files, "files", false, "Include files in addition to folders") cmd.Flags().BoolVarP(&jsonOutput, "json", "j", false, "Output results as JSON") + cmd.Flags().BoolVar(&myDrive, "my-drive", false, "Show My Drive only (default)") + cmd.Flags().StringVar(&driveFlag, "drive", "", "Show tree from specific shared drive (name or ID)") return cmd } // buildTree recursively builds the folder tree structure func buildTree(client drive.DriveClientInterface, folderID string, depth int, includeFiles bool) (*TreeNode, error) { + return buildTreeWithScope(client, folderID, "", depth, includeFiles) +} + +// buildTreeWithScope builds folder tree with optional root name override +func buildTreeWithScope(client drive.DriveClientInterface, folderID, rootName string, depth int, includeFiles bool) (*TreeNode, error) { // Get folder info var folderName string var folderType string @@ -78,6 +107,9 @@ func buildTree(client drive.DriveClientInterface, folderID string, depth int, in if folderID == "root" { folderName = "My Drive" folderType = "Folder" + } else if rootName != "" && depth == 2 { // First call with override + folderName = rootName + folderType = "Shared Drive" } else { folder, err := client.GetFile(folderID) if err != nil { @@ -98,13 +130,15 @@ func buildTree(client drive.DriveClientInterface, folderID string, depth int, in return node, nil } - // Build query to list children + // Build query to list children - use scope for shared drive support query := fmt.Sprintf("'%s' in parents and trashed = false", folderID) if !includeFiles { query += fmt.Sprintf(" and mimeType = '%s'", drive.MimeTypeFolder) } - children, err := client.ListFiles(query, 100) + // Use ListFilesWithScope to support shared drives + scope := drive.DriveScope{AllDrives: true} + children, err := client.ListFilesWithScope(query, 100, scope) if err != nil { return nil, fmt.Errorf("failed to list children: %w", err) } @@ -122,8 +156,8 @@ func buildTree(client drive.DriveClientInterface, folderID string, depth int, in // Process children for _, child := range children { if child.MimeType == drive.MimeTypeFolder { - // Recursively build subtree for folders - childNode, err := buildTree(client, child.ID, depth-1, includeFiles) + // Recursively build subtree for folders (don't pass rootName on recursion) + childNode, err := buildTreeWithScope(client, child.ID, "", depth-1, includeFiles) if err != nil { // Log error but continue with other children continue diff --git a/internal/cmd/drive/tree_test.go b/internal/cmd/drive/tree_test.go index 4e67e7a..40c0673 100644 --- a/internal/cmd/drive/tree_test.go +++ b/internal/cmd/drive/tree_test.go @@ -248,6 +248,11 @@ func (m *mockDriveClient) ListFiles(query string, _ int64) ([]*drive.File, error return []*drive.File{}, nil } +func (m *mockDriveClient) ListFilesWithScope(query string, pageSize int64, _ drive.DriveScope) ([]*drive.File, error) { + // Delegate to ListFiles for testing purposes + return m.ListFiles(query, pageSize) +} + func (m *mockDriveClient) DownloadFile(_ string) ([]byte, error) { return nil, fmt.Errorf("not implemented") } @@ -256,6 +261,10 @@ func (m *mockDriveClient) ExportFile(_ string, _ string) ([]byte, error) { return nil, fmt.Errorf("not implemented") } +func (m *mockDriveClient) ListSharedDrives(_ int64) ([]*drive.SharedDrive, error) { + return nil, fmt.Errorf("not implemented") +} + func TestBuildTree(t *testing.T) { t.Run("builds tree for root folder", func(t *testing.T) { mock := newMockDriveClient() diff --git a/internal/drive/client.go b/internal/drive/client.go index 92824cc..53a8871 100644 --- a/internal/drive/client.go +++ b/internal/drive/client.go @@ -34,9 +34,9 @@ func NewClient(ctx context.Context) (*Client, error) { } // fileFields defines the fields to request from the Drive API -const fileFields = "id,name,mimeType,size,createdTime,modifiedTime,parents,owners,webViewLink,shared" +const fileFields = "id,name,mimeType,size,createdTime,modifiedTime,parents,owners,webViewLink,shared,driveId" -// ListFiles returns files matching the query +// ListFiles returns files matching the query (searches My Drive only for backwards compatibility) func (c *Client) ListFiles(query string, pageSize int64) ([]*File, error) { call := c.service.Files.List(). Fields("files(" + fileFields + ")"). @@ -61,10 +61,51 @@ func (c *Client) ListFiles(query string, pageSize int64) ([]*File, error) { return files, nil } -// GetFile retrieves a single file by ID +// ListFilesWithScope returns files matching the query within the specified scope +func (c *Client) ListFilesWithScope(query string, pageSize int64, scope DriveScope) ([]*File, error) { + call := c.service.Files.List(). + Fields("files(" + fileFields + ")"). + OrderBy("modifiedTime desc"). + SupportsAllDrives(true). + IncludeItemsFromAllDrives(true) + + // Set corpora based on scope + if scope.DriveID != "" { + // Specific shared drive + call = call.Corpora("drive").DriveId(scope.DriveID) + } else if scope.MyDriveOnly { + // My Drive only + call = call.Corpora("user") + } else if scope.AllDrives { + // Search everywhere + call = call.Corpora("allDrives") + } + // If no scope flags set, default behavior (no corpora set) + + if query != "" { + call = call.Q(query) + } + if pageSize > 0 { + call = call.PageSize(pageSize) + } + + resp, err := call.Do() + if err != nil { + return nil, fmt.Errorf("failed to list files: %w", err) + } + + files := make([]*File, 0, len(resp.Files)) + for _, f := range resp.Files { + files = append(files, ParseFile(f)) + } + return files, nil +} + +// GetFile retrieves a single file by ID (supports files in shared drives) func (c *Client) GetFile(fileID string) (*File, error) { f, err := c.service.Files.Get(fileID). Fields(fileFields). + SupportsAllDrives(true). Do() if err != nil { return nil, fmt.Errorf("failed to get file: %w", err) @@ -74,7 +115,9 @@ func (c *Client) GetFile(fileID string) (*File, error) { // DownloadFile downloads a regular (non-Google Workspace) file func (c *Client) DownloadFile(fileID string) ([]byte, error) { - resp, err := c.service.Files.Get(fileID).Download() + resp, err := c.service.Files.Get(fileID). + SupportsAllDrives(true). + Download() if err != nil { return nil, fmt.Errorf("failed to download file: %w", err) } @@ -101,3 +144,40 @@ func (c *Client) ExportFile(fileID string, mimeType string) ([]byte, error) { } return data, nil } + +// ListSharedDrives returns all shared drives accessible to the user +func (c *Client) ListSharedDrives(pageSize int64) ([]*SharedDrive, error) { + var allDrives []*SharedDrive + pageToken := "" + + for { + call := c.service.Drives.List(). + Fields("drives(id,name),nextPageToken") + + if pageSize > 0 { + call = call.PageSize(pageSize) + } + if pageToken != "" { + call = call.PageToken(pageToken) + } + + resp, err := call.Do() + if err != nil { + return nil, fmt.Errorf("failed to list shared drives: %w", err) + } + + for _, d := range resp.Drives { + allDrives = append(allDrives, &SharedDrive{ + ID: d.Id, + Name: d.Name, + }) + } + + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + return allDrives, nil +} diff --git a/internal/drive/files.go b/internal/drive/files.go index d5356d0..b4e0a73 100644 --- a/internal/drive/files.go +++ b/internal/drive/files.go @@ -20,6 +20,20 @@ type File struct { Owners []string `json:"owners,omitempty"` WebViewLink string `json:"webViewLink,omitempty"` Shared bool `json:"shared"` + DriveID string `json:"driveId,omitempty"` // Shared drive ID if file is in a shared drive +} + +// SharedDrive represents a Google Shared Drive (formerly Team Drive) +type SharedDrive struct { + ID string `json:"id"` + Name string `json:"name"` +} + +// DriveScope defines where to search for files +type DriveScope struct { + AllDrives bool // Search everywhere (My Drive + all shared drives) + MyDriveOnly bool // Restrict to personal My Drive only + DriveID string // Specific shared drive ID } // ParseFile converts a Google Drive API File to our simplified File struct @@ -32,6 +46,7 @@ func ParseFile(f *drive.File) *File { Parents: f.Parents, WebViewLink: f.WebViewLink, Shared: f.Shared, + DriveID: f.DriveId, } // Parse timestamps diff --git a/internal/drive/interfaces.go b/internal/drive/interfaces.go index f4e36cc..8b55972 100644 --- a/internal/drive/interfaces.go +++ b/internal/drive/interfaces.go @@ -3,10 +3,13 @@ package drive // DriveClientInterface defines the interface for Drive client operations. // This enables unit testing through mock implementations. type DriveClientInterface interface { - // ListFiles returns files matching the query + // ListFiles returns files matching the query (searches My Drive only for backwards compatibility) ListFiles(query string, pageSize int64) ([]*File, error) - // GetFile retrieves a single file by ID + // ListFilesWithScope returns files matching the query within the specified scope + ListFilesWithScope(query string, pageSize int64, scope DriveScope) ([]*File, error) + + // GetFile retrieves a single file by ID (supports all drives) GetFile(fileID string) (*File, error) // DownloadFile downloads a regular (non-Google Workspace) file @@ -14,6 +17,9 @@ type DriveClientInterface interface { // ExportFile exports a Google Workspace file to the specified MIME type ExportFile(fileID string, mimeType string) ([]byte, error) + + // ListSharedDrives returns all shared drives accessible to the user + ListSharedDrives(pageSize int64) ([]*SharedDrive, error) } // Verify that Client implements DriveClientInterface diff --git a/internal/testutil/fixtures.go b/internal/testutil/fixtures.go index cb58992..84c68e2 100644 --- a/internal/testutil/fixtures.go +++ b/internal/testutil/fixtures.go @@ -238,3 +238,37 @@ func SampleGoogleDoc(id string) *driveapi.File { Shared: false, } } + +// SampleSharedDrive returns a sample shared drive for testing +func SampleSharedDrive(id, name string) *driveapi.SharedDrive { + return &driveapi.SharedDrive{ + ID: id, + Name: name, + } +} + +// SampleSharedDrives returns a set of sample shared drives for testing +func SampleSharedDrives() []*driveapi.SharedDrive { + return []*driveapi.SharedDrive{ + {ID: "0ALengineering123", Name: "Engineering"}, + {ID: "0ALmarketing456", Name: "Marketing"}, + {ID: "0ALfinance789", Name: "Finance Team"}, + } +} + +// SampleSharedDriveFile returns a file that belongs to a shared drive +func SampleSharedDriveFile(id, driveID string) *driveapi.File { + return &driveapi.File{ + ID: id, + Name: "shared-document.pdf", + MimeType: "application/pdf", + Size: 4096, + CreatedTime: time.Date(2024, 1, 10, 9, 0, 0, 0, time.UTC), + ModifiedTime: time.Date(2024, 1, 15, 14, 30, 0, 0, time.UTC), + Parents: []string{driveID}, + Owners: []string{}, + WebViewLink: "https://drive.google.com/file/d/" + id, + Shared: true, + DriveID: driveID, + } +} diff --git a/internal/testutil/mocks.go b/internal/testutil/mocks.go index b8827bb..4a10737 100644 --- a/internal/testutil/mocks.go +++ b/internal/testutil/mocks.go @@ -173,10 +173,12 @@ func (m *MockContactsClient) ListContactGroups(pageToken string, pageSize int64) // MockDriveClient is a configurable mock for DriveClientInterface. type MockDriveClient struct { - ListFilesFunc func(query string, pageSize int64) ([]*driveapi.File, error) - GetFileFunc func(fileID string) (*driveapi.File, error) - DownloadFileFunc func(fileID string) ([]byte, error) - ExportFileFunc func(fileID, mimeType string) ([]byte, error) + ListFilesFunc func(query string, pageSize int64) ([]*driveapi.File, error) + ListFilesWithScopeFunc func(query string, pageSize int64, scope driveapi.DriveScope) ([]*driveapi.File, error) + GetFileFunc func(fileID string) (*driveapi.File, error) + DownloadFileFunc func(fileID string) ([]byte, error) + ExportFileFunc func(fileID, mimeType string) ([]byte, error) + ListSharedDrivesFunc func(pageSize int64) ([]*driveapi.SharedDrive, error) } // Verify MockDriveClient implements DriveClientInterface @@ -189,6 +191,17 @@ func (m *MockDriveClient) ListFiles(query string, pageSize int64) ([]*driveapi.F return nil, nil } +func (m *MockDriveClient) ListFilesWithScope(query string, pageSize int64, scope driveapi.DriveScope) ([]*driveapi.File, error) { + if m.ListFilesWithScopeFunc != nil { + return m.ListFilesWithScopeFunc(query, pageSize, scope) + } + // Fall back to ListFiles if no scope function defined + if m.ListFilesFunc != nil { + return m.ListFilesFunc(query, pageSize) + } + return nil, nil +} + func (m *MockDriveClient) GetFile(fileID string) (*driveapi.File, error) { if m.GetFileFunc != nil { return m.GetFileFunc(fileID) @@ -209,3 +222,10 @@ func (m *MockDriveClient) ExportFile(fileID, mimeType string) ([]byte, error) { } return nil, nil } + +func (m *MockDriveClient) ListSharedDrives(pageSize int64) ([]*driveapi.SharedDrive, error) { + if m.ListSharedDrivesFunc != nil { + return m.ListSharedDrivesFunc(pageSize) + } + return nil, nil +}