From 50fab66d3b0d51d216b9626ace9ea46933d92e45 Mon Sep 17 00:00:00 2001
From: LJ <linghua@cocoindex.io>
Date: Tue, 18 Mar 2025 16:37:58 -0700
Subject: [PATCH] Support multiple root folders for a `GoogleDrive` data
 source.

---
 python/cocoindex/sources.py     |  2 +-
 src/ops/sources/google_drive.rs | 12 +++++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/python/cocoindex/sources.py b/python/cocoindex/sources.py
index decaeca1..12bac873 100644
--- a/python/cocoindex/sources.py
+++ b/python/cocoindex/sources.py
@@ -24,5 +24,5 @@ class GoogleDrive(op.SourceSpec):
     _op_category = op.OpCategory.SOURCE
 
     service_account_credential_path: str
-    root_folder_id: str
+    root_folder_ids: list[str]
     binary: bool = False
diff --git a/src/ops/sources/google_drive.rs b/src/ops/sources/google_drive.rs
index e07995dd..912bf933 100644
--- a/src/ops/sources/google_drive.rs
+++ b/src/ops/sources/google_drive.rs
@@ -73,13 +73,13 @@ fn is_supported_file_type(mime_type: &str) -> bool {
 pub struct Spec {
     service_account_credential_path: String,
     binary: bool,
-    root_folder_id: String,
+    root_folder_ids: Vec<String>,
 }
 
 struct Executor {
     drive_hub: DriveHub<HttpsConnector<HttpConnector>>,
     binary: bool,
-    root_folder_id: String,
+    root_folder_ids: Vec<String>,
 }
 
 impl Executor {
@@ -102,7 +102,7 @@ impl Executor {
         Ok(Self {
             drive_hub,
             binary: spec.binary,
-            root_folder_id: spec.root_folder_id,
+            root_folder_ids: spec.root_folder_ids,
         })
     }
 }
@@ -176,8 +176,10 @@ impl Executor {
 impl SourceExecutor for Executor {
     async fn list_keys(&self) -> Result<Vec<KeyValue>> {
         let mut result = IndexSet::new();
-        self.traverse_folder(&self.root_folder_id, &mut IndexSet::new(), &mut result)
-            .await?;
+        for root_folder_id in &self.root_folder_ids {
+            self.traverse_folder(root_folder_id, &mut IndexSet::new(), &mut result)
+                .await?;
+        }
         Ok(result.into_iter().collect())
     }