Skip to content

Commit 871b527

Browse files
authored
Support multiple root folders for a GoogleDrive data source. (#163)
1 parent db11fd8 commit 871b527

File tree

2 files changed

+8
-6
lines changed

2 files changed

+8
-6
lines changed

python/cocoindex/sources.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,5 +24,5 @@ class GoogleDrive(op.SourceSpec):
2424
_op_category = op.OpCategory.SOURCE
2525

2626
service_account_credential_path: str
27-
root_folder_id: str
27+
root_folder_ids: list[str]
2828
binary: bool = False

src/ops/sources/google_drive.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -73,13 +73,13 @@ fn is_supported_file_type(mime_type: &str) -> bool {
7373
pub struct Spec {
7474
service_account_credential_path: String,
7575
binary: bool,
76-
root_folder_id: String,
76+
root_folder_ids: Vec<String>,
7777
}
7878

7979
struct Executor {
8080
drive_hub: DriveHub<HttpsConnector<HttpConnector>>,
8181
binary: bool,
82-
root_folder_id: String,
82+
root_folder_ids: Vec<String>,
8383
}
8484

8585
impl Executor {
@@ -102,7 +102,7 @@ impl Executor {
102102
Ok(Self {
103103
drive_hub,
104104
binary: spec.binary,
105-
root_folder_id: spec.root_folder_id,
105+
root_folder_ids: spec.root_folder_ids,
106106
})
107107
}
108108
}
@@ -176,8 +176,10 @@ impl Executor {
176176
impl SourceExecutor for Executor {
177177
async fn list_keys(&self) -> Result<Vec<KeyValue>> {
178178
let mut result = IndexSet::new();
179-
self.traverse_folder(&self.root_folder_id, &mut IndexSet::new(), &mut result)
180-
.await?;
179+
for root_folder_id in &self.root_folder_ids {
180+
self.traverse_folder(root_folder_id, &mut IndexSet::new(), &mut result)
181+
.await?;
182+
}
181183
Ok(result.into_iter().collect())
182184
}
183185

0 commit comments

Comments
 (0)