dataprof_db/connectors/
sqlite.rs

1//! SQLite database connector (embedded database)
2
3#[cfg(not(feature = "sqlite"))]
4use super::common::feature_not_enabled_error;
5#[cfg(feature = "sqlite")]
6use super::common::{build_count_query, not_connected_error};
7use crate::connection::ConnectionInfo;
8#[cfg(feature = "sqlite")]
9use crate::security::validate_sql_identifier;
10use crate::{DataProfilerError, DatabaseConfig, DatabaseConnector, QueryColumns};
11#[cfg(feature = "sqlite")]
12use crate::{process_rows_to_columns, streaming_profile_loop};
13use async_trait::async_trait;
14
15#[cfg(feature = "sqlite")]
16use {sqlx::sqlite::SqlitePool, sqlx::sqlite::SqlitePoolOptions};
17
18/// SQLite embedded database connector
19pub struct SqliteConnector {
20    #[allow(dead_code)]
21    config: DatabaseConfig,
22    #[allow(dead_code)]
23    connection_info: ConnectionInfo,
24    #[cfg(feature = "sqlite")]
25    pool: Option<SqlitePool>,
26    #[cfg(not(feature = "sqlite"))]
27    #[allow(dead_code)]
28    pool: Option<()>,
29}
30
31impl SqliteConnector {
32    /// Create a new SQLite connector
33    pub fn new(config: DatabaseConfig) -> Result<Self, DataProfilerError> {
34        let connection_info = ConnectionInfo::parse(&config.connection_string)?;
35
36        if connection_info.database_type() != "sqlite" {
37            return Err(DataProfilerError::DatabaseConfigError {
38                message: format!(
39                    "Invalid connection string for SQLite: {}",
40                    config.connection_string
41                ),
42            });
43        }
44
45        Ok(Self {
46            config,
47            connection_info,
48            pool: None,
49        })
50    }
51
52    /// Get the database file path
53    #[allow(dead_code)]
54    fn get_db_path(&self) -> Result<String, DataProfilerError> {
55        if let Some(path) = &self.connection_info.path {
56            Ok(path.clone())
57        } else if let Some(db) = &self.connection_info.database {
58            Ok(db.clone())
59        } else {
60            Err(DataProfilerError::DatabaseConfigError {
61                message: "No database path specified for SQLite".to_string(),
62            })
63        }
64    }
65}
66
67#[async_trait]
68impl DatabaseConnector for SqliteConnector {
69    async fn connect(&mut self) -> Result<(), DataProfilerError> {
70        #[cfg(feature = "sqlite")]
71        {
72            let db_path = self.get_db_path()?;
73            let connection_string = if db_path == ":memory:" {
74                "sqlite::memory:".to_string()
75            } else {
76                format!("sqlite://{}", db_path)
77            };
78
79            let pool = SqlitePoolOptions::new()
80                .max_connections(1)
81                .acquire_timeout(
82                    self.config
83                        .connection_timeout
84                        .unwrap_or(std::time::Duration::from_secs(30)),
85                )
86                .connect(&connection_string)
87                .await
88                .map_err(|e| {
89                    DataProfilerError::database_connection(&format!(
90                        "Failed to connect to SQLite: {}",
91                        e
92                    ))
93                })?;
94
95            self.pool = Some(pool);
96            Ok(())
97        }
98
99        #[cfg(not(feature = "sqlite"))]
100        {
101            Err(DataProfilerError::database_feature_disabled(
102                "SQLite", "sqlite",
103            ))
104        }
105    }
106
107    async fn disconnect(&mut self) -> Result<(), DataProfilerError> {
108        #[cfg(feature = "sqlite")]
109        {
110            if let Some(pool) = &self.pool {
111                pool.close().await;
112                self.pool = None;
113            }
114        }
115        Ok(())
116    }
117
118    #[allow(unused_variables)]
119    async fn profile_query(&mut self, query: &str) -> Result<QueryColumns, DataProfilerError> {
120        #[cfg(feature = "sqlite")]
121        {
122            let pool = self.pool.as_ref().ok_or_else(not_connected_error)?;
123
124            let rows = sqlx::query(query).fetch_all(pool).await.map_err(|e| {
125                DataProfilerError::database_query(&format!("Query execution failed: {}", e))
126            })?;
127
128            process_rows_to_columns!(rows)
129        }
130
131        #[cfg(not(feature = "sqlite"))]
132        Err(feature_not_enabled_error("SQLite", "sqlite"))
133    }
134
135    #[allow(unused_variables)]
136    async fn profile_query_streaming(
137        &mut self,
138        query: &str,
139        batch_size: usize,
140    ) -> Result<QueryColumns, DataProfilerError> {
141        #[cfg(feature = "sqlite")]
142        {
143            let pool = self.pool.as_ref().ok_or_else(not_connected_error)?;
144
145            let count_query = build_count_query(query)?;
146            let total_rows: i64 = sqlx::query_scalar(&count_query)
147                .fetch_one(pool)
148                .await
149                .map_err(|e| {
150                    DataProfilerError::database_query(&format!("Failed to count rows: {}", e))
151                })?;
152
153            streaming_profile_loop!(pool, query, batch_size, total_rows, "SQLite")
154        }
155
156        #[cfg(not(feature = "sqlite"))]
157        Err(feature_not_enabled_error("SQLite", "sqlite"))
158    }
159
160    #[allow(unused_variables)]
161    async fn get_table_schema(
162        &mut self,
163        table_name: &str,
164    ) -> Result<Vec<String>, DataProfilerError> {
165        #[cfg(feature = "sqlite")]
166        {
167            use sqlx::Row;
168
169            let pool = self.pool.as_ref().ok_or_else(not_connected_error)?;
170
171            let query = format!("PRAGMA table_info({})", table_name);
172
173            let rows = sqlx::query(&query).fetch_all(pool).await.map_err(|e| {
174                DataProfilerError::database_query(&format!("Failed to get table schema: {}", e))
175            })?;
176
177            let mut columns = Vec::new();
178            for row in rows {
179                let column_name: String = row.try_get(1).map_err(|e| {
180                    DataProfilerError::database_query(&format!("Failed to read column name: {}", e))
181                })?;
182                columns.push(column_name);
183            }
184
185            Ok(columns)
186        }
187
188        #[cfg(not(feature = "sqlite"))]
189        Err(feature_not_enabled_error("SQLite", "sqlite"))
190    }
191
192    #[allow(unused_variables)]
193    async fn count_table_rows(&mut self, table_name: &str) -> Result<u64, DataProfilerError> {
194        #[cfg(feature = "sqlite")]
195        {
196            let pool = self.pool.as_ref().ok_or_else(not_connected_error)?;
197
198            validate_sql_identifier(table_name)?;
199            let query = format!("SELECT COUNT(*) FROM {}", table_name);
200            let count: i64 = sqlx::query_scalar(&query)
201                .fetch_one(pool)
202                .await
203                .map_err(|e| {
204                    DataProfilerError::database_query(&format!("Failed to count rows: {}", e))
205                })?;
206
207            Ok(count as u64)
208        }
209
210        #[cfg(not(feature = "sqlite"))]
211        Err(feature_not_enabled_error("SQLite", "sqlite"))
212    }
213
214    async fn test_connection(&mut self) -> Result<bool, DataProfilerError> {
215        #[cfg(feature = "sqlite")]
216        {
217            let pool = self.pool.as_ref().ok_or_else(not_connected_error)?;
218
219            let result: i32 = sqlx::query_scalar("SELECT 1")
220                .fetch_one(pool)
221                .await
222                .map_err(|e| {
223                    DataProfilerError::database_query(&format!("Connection test failed: {}", e))
224                })?;
225
226            Ok(result == 1)
227        }
228
229        #[cfg(not(feature = "sqlite"))]
230        Err(feature_not_enabled_error("SQLite", "sqlite"))
231    }
232}
dataprof_db/connectors/sqlite.rs

dataprof_db/connectors/
sqlite.rs
