Skip to main content

sqlite_vtable_opendal/backends/
gdrive.rs

1//! Google Drive storage backend implementation
2//!
3//! This backend allows querying Google Drive files and folders using SQL.
4//! Requires a Google Drive access token for authentication.
5
6use crate::backends::StorageBackend;
7use crate::error::{Result, VTableError};
8use crate::types::{FileMetadata, QueryConfig};
9use async_trait::async_trait;
10use futures_util::TryStreamExt;
11use opendal::{services::Gdrive, EntryMode, Metakey, Operator};
12use std::path::Path;
13
14/// Google Drive storage backend
15///
16/// This backend uses OpenDAL's Gdrive service to list files from Google Drive.
17///
18/// # Authentication
19///
20/// Requires a Google Drive access token. You can obtain one from:
21/// https://console.cloud.google.com/apis/credentials
22///
23/// # Example
24///
25/// ```rust,ignore
26/// use sqlite_vtable_opendal::backends::gdrive::GdriveBackend;
27/// use sqlite_vtable_opendal::types::QueryConfig;
28///
29/// let backend = GdriveBackend::new("YOUR_ACCESS_TOKEN", "/");
30/// let config = QueryConfig::default();
31/// let files = backend.list_files(&config).await?;
32/// ```
33pub struct GdriveBackend {
34    /// Google Drive access token
35    access_token: String,
36    /// Base path in Google Drive (e.g., "/" for root)
37    base_path: String,
38}
39
40impl GdriveBackend {
41    /// Create a new Google Drive backend
42    ///
43    /// # Arguments
44    ///
45    /// * `access_token` - Google Drive API access token
46    /// * `base_path` - Base path to query from (e.g., "/" for root)
47    ///
48    /// # Example
49    ///
50    /// ```
51    /// use sqlite_vtable_opendal::backends::gdrive::GdriveBackend;
52    ///
53    /// let backend = GdriveBackend::new("token", "/My Documents");
54    /// ```
55    pub fn new(access_token: impl Into<String>, base_path: impl Into<String>) -> Self {
56        Self {
57            access_token: access_token.into(),
58            base_path: base_path.into(),
59        }
60    }
61
62    /// Create an OpenDAL operator for Google Drive
63    fn create_operator(&self) -> Result<Operator> {
64        let builder = Gdrive::default()
65            .access_token(&self.access_token)
66            .root(&self.base_path);
67
68        Operator::new(builder)
69            .map(|op| op.finish())
70            .map_err(|e| VTableError::OpenDal(e))
71    }
72}
73
74#[async_trait]
75impl StorageBackend for GdriveBackend {
76    async fn list_files(&self, config: &QueryConfig) -> Result<Vec<FileMetadata>> {
77        let operator = self.create_operator()?;
78        let mut results = Vec::new();
79
80        // Normalize the path
81        let normalized_path = if config.root_path.is_empty() || config.root_path == "/" {
82            "".to_string()
83        } else {
84            let clean_path = config.root_path.trim_matches('/');
85            if clean_path.is_empty() {
86                "".to_string()
87            } else {
88                format!("/{}", clean_path)
89            }
90        };
91
92        // Create lister with metadata keys
93        let lister_builder = operator.lister_with(&normalized_path);
94
95        let mut lister = lister_builder
96            .recursive(config.recursive)
97            .metakey(
98                Metakey::ContentLength
99                    | Metakey::ContentMd5
100                    | Metakey::ContentType
101                    | Metakey::Mode
102                    | Metakey::LastModified,
103            )
104            .await
105            .map_err(|e| VTableError::OpenDal(e))?;
106
107        // Iterate through entries
108        while let Some(entry) = lister.try_next().await.map_err(|e| VTableError::OpenDal(e))? {
109            let entry_path = entry.path();
110            let entry_mode = entry.metadata().mode();
111
112            // Skip the root directory itself
113            if entry_path.is_empty() || entry_path == "/" || entry_path == "." {
114                continue;
115            }
116
117            let full_path = if entry_path.starts_with('/') {
118                entry_path.to_string()
119            } else {
120                format!("/{}", entry_path)
121            };
122
123            // Extract file name from path
124            let name = Path::new(&full_path)
125                .file_name()
126                .map(|s| s.to_string_lossy().to_string())
127                .unwrap_or_else(|| {
128                    let clean_path = entry_path.trim_end_matches('/');
129                    Path::new(clean_path)
130                        .file_name()
131                        .map(|s| s.to_string_lossy().to_string())
132                        .unwrap_or_default()
133                });
134
135            if entry_mode == EntryMode::FILE {
136                // Fetch detailed metadata for files
137                let metadata = operator
138                    .stat(&full_path)
139                    .await
140                    .map_err(|e| VTableError::OpenDal(e))?;
141
142                // Optionally fetch content
143                let content = if config.fetch_content {
144                    operator
145                        .read(&full_path)
146                        .await
147                        .ok()
148                        .map(|bytes| bytes.to_vec())
149                } else {
150                    None
151                };
152
153                results.push(FileMetadata {
154                    name,
155                    path: full_path.clone(),
156                    size: metadata.content_length(),
157                    last_modified: metadata.last_modified().map(|dt| dt.to_string()),
158                    etag: metadata.content_md5().map(|md5| md5.to_string()),
159                    is_dir: false,
160                    content_type: Path::new(&full_path)
161                        .extension()
162                        .and_then(|ext| ext.to_str())
163                        .map(|ext| ext.to_string()),
164                    content,
165                });
166
167                // Apply limit if specified
168                if let Some(limit) = config.limit {
169                    if results.len() >= limit + config.offset {
170                        break;
171                    }
172                }
173            } else if entry_mode == EntryMode::DIR {
174                // Add directory entry
175                results.push(FileMetadata {
176                    name,
177                    path: full_path,
178                    size: 0,
179                    last_modified: None,
180                    etag: None,
181                    is_dir: true,
182                    content_type: Some("directory".to_string()),
183                    content: None,
184                });
185
186                // Apply limit if specified
187                if let Some(limit) = config.limit {
188                    if results.len() >= limit + config.offset {
189                        break;
190                    }
191                }
192            }
193        }
194
195        // Apply offset
196        if config.offset > 0 && config.offset < results.len() {
197            results = results.into_iter().skip(config.offset).collect();
198        }
199
200        Ok(results)
201    }
202
203    fn backend_name(&self) -> &'static str {
204        "gdrive"
205    }
206}
207
208/// Register the gdrive virtual table with SQLite
209///
210/// This function creates a virtual table module that allows querying
211/// Google Drive files using SQL.
212///
213/// # Arguments
214///
215/// * `conn` - SQLite connection
216/// * `module_name` - Name for the virtual table (e.g., "gdrive_files")
217/// * `access_token` - Google Drive access token
218/// * `base_path` - Base path in Google Drive (e.g., "/" for root)
219///
220/// # Example
221///
222/// ```rust,ignore
223/// use rusqlite::Connection;
224/// use sqlite_vtable_opendal::backends::gdrive;
225///
226/// let conn = Connection::open_in_memory()?;
227/// gdrive::register(&conn, "gdrive_files", "YOUR_TOKEN", "/")?;
228///
229/// // Now you can query: SELECT * FROM gdrive_files
230/// ```
231pub fn register(
232    conn: &rusqlite::Connection,
233    module_name: &str,
234    access_token: impl Into<String>,
235    base_path: impl Into<String>,
236) -> rusqlite::Result<()> {
237    use crate::types::{columns, QueryConfig};
238    use rusqlite::{
239        ffi,
240        vtab::{self, eponymous_only_module, IndexInfo, VTab, VTabCursor, VTabKind},
241    };
242    use std::os::raw::c_int;
243
244    let token = access_token.into();
245    let path = base_path.into();
246
247    // Create a specific table type for Google Drive
248    #[repr(C)]
249    struct GdriveTable {
250        base: ffi::sqlite3_vtab,
251        access_token: String,
252        base_path: String,
253    }
254
255    // Create a specific cursor type for Google Drive
256    #[repr(C)]
257    struct GdriveCursor {
258        base: ffi::sqlite3_vtab_cursor,
259        files: Vec<crate::types::FileMetadata>,
260        current_row: usize,
261        access_token: String,
262        base_path: String,
263    }
264
265    impl GdriveCursor {
266        fn new(access_token: String, base_path: String) -> Self {
267            Self {
268                base: ffi::sqlite3_vtab_cursor::default(),
269                files: Vec::new(),
270                current_row: 0,
271                access_token,
272                base_path,
273            }
274        }
275    }
276
277    unsafe impl VTabCursor for GdriveCursor {
278        fn filter(
279            &mut self,
280            _idx_num: c_int,
281            _idx_str: Option<&str>,
282            _args: &vtab::Values<'_>,
283        ) -> rusqlite::Result<()> {
284            // Create backend and fetch files
285            let backend = GdriveBackend::new(&self.access_token, &self.base_path);
286            let config = QueryConfig::default();
287
288            // Fetch files from the backend (blocking the async call)
289            let files = tokio::task::block_in_place(|| {
290                tokio::runtime::Handle::current().block_on(async {
291                    backend.list_files(&config).await
292                })
293            })
294            .map_err(|e| rusqlite::Error::ModuleError(e.to_string()))?;
295
296            self.files = files;
297            self.current_row = 0;
298            Ok(())
299        }
300
301        fn next(&mut self) -> rusqlite::Result<()> {
302            self.current_row += 1;
303            Ok(())
304        }
305
306        fn eof(&self) -> bool {
307            self.current_row >= self.files.len()
308        }
309
310        fn column(&self, ctx: &mut vtab::Context, col_index: c_int) -> rusqlite::Result<()> {
311            if self.current_row >= self.files.len() {
312                return Ok(());
313            }
314
315            let file = &self.files[self.current_row];
316
317            match col_index {
318                columns::PATH => ctx.set_result(&file.path),
319                columns::SIZE => ctx.set_result(&(file.size as i64)),
320                columns::LAST_MODIFIED => ctx.set_result(&file.last_modified),
321                columns::ETAG => ctx.set_result(&file.etag),
322                columns::IS_DIR => ctx.set_result(&file.is_dir),
323                columns::CONTENT_TYPE => ctx.set_result(&file.content_type),
324                columns::NAME => ctx.set_result(&file.name),
325                columns::CONTENT => {
326                    if let Some(ref content) = file.content {
327                        ctx.set_result(&content.as_slice())
328                    } else {
329                        ctx.set_result::<Option<&[u8]>>(&None)
330                    }
331                }
332                _ => Ok(()),
333            }
334        }
335
336        fn rowid(&self) -> rusqlite::Result<i64> {
337            Ok(self.current_row as i64)
338        }
339    }
340
341    impl vtab::CreateVTab<'_> for GdriveTable {
342        const KIND: VTabKind = VTabKind::EponymousOnly;
343    }
344
345    unsafe impl VTab<'_> for GdriveTable {
346        type Aux = (String, String);
347        type Cursor = GdriveCursor;
348
349        fn connect(
350            _db: &mut vtab::VTabConnection,
351            aux: Option<&Self::Aux>,
352            _args: &[&[u8]],
353        ) -> rusqlite::Result<(String, Self)> {
354            let schema = "
355                CREATE TABLE x(
356                    path TEXT,
357                    size INTEGER,
358                    last_modified TEXT,
359                    etag TEXT,
360                    is_dir INTEGER,
361                    content_type TEXT,
362                    name TEXT,
363                    content BLOB
364                )
365            ";
366
367            let (access_token, base_path) = if let Some((token, path)) = aux {
368                (token.clone(), path.clone())
369            } else {
370                ("/".to_string(), "/".to_string())
371            };
372
373            Ok((
374                schema.to_owned(),
375                GdriveTable {
376                    base: ffi::sqlite3_vtab::default(),
377                    access_token,
378                    base_path,
379                },
380            ))
381        }
382
383        fn best_index(&self, info: &mut IndexInfo) -> rusqlite::Result<()> {
384            info.set_estimated_cost(1000.0);
385            Ok(())
386        }
387
388        fn open(&mut self) -> rusqlite::Result<Self::Cursor> {
389            Ok(GdriveCursor::new(
390                self.access_token.clone(),
391                self.base_path.clone(),
392            ))
393        }
394    }
395
396    conn.create_module(
397        module_name,
398        eponymous_only_module::<GdriveTable>(),
399        Some((token, path)),
400    )
401}
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406
407    #[test]
408    fn test_backend_creation() {
409        let backend = GdriveBackend::new("test_token", "/My Documents");
410        assert_eq!(backend.access_token, "test_token");
411        assert_eq!(backend.base_path, "/My Documents");
412        assert_eq!(backend.backend_name(), "gdrive");
413    }
414
415    #[test]
416    fn test_backend_with_root_path() {
417        let backend = GdriveBackend::new("token", "/");
418        assert_eq!(backend.base_path, "/");
419    }
420
421    // Note: Integration tests with actual Google Drive API would require credentials
422    // and are better suited for manual testing or CI with secrets
423}