chdb_rust/session.rs
1//! Session management for persistent chDB databases.
2//!
3//! This module provides the [`Session`] and [`SessionBuilder`] types for managing
4//! persistent database connections with automatic cleanup.
5
6use std::fs;
7use std::path::PathBuf;
8
9use crate::arg::Arg;
10use crate::connection::Connection;
11use crate::error::Error;
12use crate::format::OutputFormat;
13use crate::query_result::QueryResult;
14
15/// Builder for creating [`Session`] instances.
16///
17/// `SessionBuilder` provides a fluent API for configuring and creating sessions.
18/// Use [`new`](Self::new) to create a new builder, configure it with the desired
19/// options, and call [`build`](Self::build) to create the session.
20///
21/// # Examples
22///
23/// ```no_run
24/// use chdb_rust::session::SessionBuilder;
25///
26/// // Create a session with default settings
27/// let session = SessionBuilder::new()
28/// .with_data_path("/tmp/mydb")
29/// .with_auto_cleanup(true)
30/// .build()?;
31/// # Ok::<(), chdb_rust::error::Error>(())
32/// ```
33pub struct SessionBuilder<'a> {
34 data_path: PathBuf,
35 default_format: OutputFormat,
36 _marker: std::marker::PhantomData<&'a ()>,
37 auto_cleanup: bool,
38}
39
40/// A session representing a persistent connection to a chDB database.
41///
42/// A `Session` manages a connection to a persistent database stored on disk.
43/// Unlike stateless queries, sessions allow you to create tables, insert data,
44/// and maintain state across multiple queries.
45///
46/// # Thread Safety
47///
48/// `Session` contains a [`Connection`] which implements
49/// `Send`, so sessions can be safely transferred between threads. However, concurrent
50/// access to the same session should be synchronized.
51///
52/// # Examples
53///
54/// ```no_run
55/// use chdb_rust::session::SessionBuilder;
56/// use chdb_rust::arg::Arg;
57/// use chdb_rust::format::OutputFormat;
58///
59/// let session = SessionBuilder::new()
60/// .with_data_path("/tmp/mydb")
61/// .with_auto_cleanup(true)
62/// .build()?;
63///
64/// // Create a table
65/// session.execute(
66/// "CREATE TABLE users (id UInt64, name String) ENGINE = MergeTree() ORDER BY id",
67/// None
68/// )?;
69///
70/// // Insert data
71/// session.execute("INSERT INTO users VALUES (1, 'Alice')", None)?;
72///
73/// // Query data
74/// let result = session.execute(
75/// "SELECT * FROM users",
76/// Some(&[Arg::OutputFormat(OutputFormat::JSONEachRow)])
77/// )?;
78/// println!("{}", result.data_utf8_lossy());
79/// # Ok::<(), chdb_rust::error::Error>(())
80/// ```
81#[derive(Debug)]
82pub struct Session {
83 conn: Connection,
84 data_path: String,
85 default_format: OutputFormat,
86 auto_cleanup: bool,
87}
88
89impl<'a> SessionBuilder<'a> {
90 /// Create a new `SessionBuilder` with default settings.
91 ///
92 /// The default settings are:
93 /// - Data path: `./chdb` in the current working directory
94 /// - Output format: `TabSeparated`
95 /// - Auto cleanup: `false`
96 ///
97 /// # Examples
98 ///
99 /// ```no_run
100 /// use chdb_rust::session::SessionBuilder;
101 ///
102 /// let builder = SessionBuilder::new();
103 /// # Ok::<(), chdb_rust::error::Error>(())
104 /// ```
105 pub fn new() -> Self {
106 let mut data_path = std::env::current_dir().unwrap();
107 data_path.push("chdb");
108
109 Self {
110 data_path,
111 default_format: OutputFormat::TabSeparated,
112 _marker: std::marker::PhantomData,
113 auto_cleanup: false,
114 }
115 }
116
117 /// Set the data path for the session.
118 ///
119 /// This specifies the filesystem path where the database will be stored.
120 /// The directory will be created if it doesn't exist.
121 ///
122 /// # Arguments
123 ///
124 /// * `path` - The path where the database should be stored
125 ///
126 /// # Examples
127 ///
128 /// ```no_run
129 /// use chdb_rust::session::SessionBuilder;
130 ///
131 /// let builder = SessionBuilder::new()
132 /// .with_data_path("/tmp/mydb");
133 /// # Ok::<(), chdb_rust::error::Error>(())
134 /// ```
135 pub fn with_data_path(mut self, path: impl Into<PathBuf>) -> Self {
136 self.data_path = path.into();
137 self
138 }
139
140 /// Add a query argument to the session builder.
141 ///
142 /// Currently, only `OutputFormat` arguments are supported and will be used
143 /// as the default output format for queries executed on this session.
144 ///
145 /// # Arguments
146 ///
147 /// * `arg` - The argument to add (currently only `OutputFormat` is supported)
148 ///
149 /// # Examples
150 ///
151 /// ```no_run
152 /// use chdb_rust::session::SessionBuilder;
153 /// use chdb_rust::arg::Arg;
154 /// use chdb_rust::format::OutputFormat;
155 ///
156 /// let builder = SessionBuilder::new()
157 /// .with_arg(Arg::OutputFormat(OutputFormat::JSONEachRow));
158 /// # Ok::<(), chdb_rust::error::Error>(())
159 /// ```
160 pub fn with_arg(mut self, arg: Arg<'a>) -> Self {
161 // Only OutputFormat is supported with the new API
162 if let Some(fmt) = arg.as_output_format() {
163 self.default_format = fmt;
164 }
165 self
166 }
167
168 /// Enable or disable automatic cleanup of the data directory.
169 ///
170 /// If set to `true`, the session will automatically delete the data directory
171 /// when it is dropped. This is useful for temporary databases.
172 ///
173 /// # Arguments
174 ///
175 /// * `value` - Whether to enable automatic cleanup
176 ///
177 /// # Examples
178 ///
179 /// ```no_run
180 /// use chdb_rust::session::SessionBuilder;
181 ///
182 /// // Session will clean up data directory on drop
183 /// let session = SessionBuilder::new()
184 /// .with_data_path("/tmp/tempdb")
185 /// .with_auto_cleanup(true)
186 /// .build()?;
187 /// # Ok::<(), chdb_rust::error::Error>(())
188 /// ```
189 pub fn with_auto_cleanup(mut self, value: bool) -> Self {
190 self.auto_cleanup = value;
191 self
192 }
193
194 /// Build the session with the configured settings.
195 ///
196 /// This creates the data directory if it doesn't exist and establishes
197 /// a connection to the database.
198 ///
199 /// # Returns
200 ///
201 /// Returns a [`Session`] if successful, or an [`Error`] if
202 /// the session cannot be created.
203 ///
204 /// # Errors
205 ///
206 /// Returns an error if:
207 /// - The data path cannot be created
208 /// - The data path has insufficient permissions
209 /// - The connection cannot be established
210 ///
211 /// # Examples
212 ///
213 /// ```no_run
214 /// use chdb_rust::session::SessionBuilder;
215 ///
216 /// let session = SessionBuilder::new()
217 /// .with_data_path("/tmp/mydb")
218 /// .build()?;
219 /// # Ok::<(), chdb_rust::error::Error>(())
220 /// ```
221 pub fn build(self) -> Result<Session, Error> {
222 let data_path = self.data_path.to_str().ok_or(Error::PathError)?.to_string();
223
224 fs::create_dir_all(&self.data_path)?;
225 if fs::metadata(&self.data_path)?.permissions().readonly() {
226 return Err(Error::InsufficientPermissions);
227 }
228
229 let conn = Connection::open_with_path(&data_path)?;
230
231 Ok(Session {
232 conn,
233 data_path,
234 default_format: self.default_format,
235 auto_cleanup: self.auto_cleanup,
236 })
237 }
238}
239
240impl Default for SessionBuilder<'_> {
241 fn default() -> Self {
242 Self::new()
243 }
244}
245
246impl Session {
247 /// Execute a query on this session.
248 ///
249 /// This executes a SQL query against the database associated with this session.
250 /// The query can create tables, insert data, or query existing data.
251 ///
252 /// # Arguments
253 ///
254 /// * `query` - The SQL query string to execute
255 /// * `query_args` - Optional array of query arguments (e.g., output format)
256 ///
257 /// # Returns
258 ///
259 /// Returns a [`QueryResult`] containing the query output,
260 /// or an [`Error`] if the query fails.
261 ///
262 /// # Examples
263 ///
264 /// ```no_run
265 /// use chdb_rust::session::SessionBuilder;
266 /// use chdb_rust::arg::Arg;
267 /// use chdb_rust::format::OutputFormat;
268 ///
269 /// let session = SessionBuilder::new()
270 /// .with_data_path("/tmp/mydb")
271 /// .with_auto_cleanup(true)
272 /// .build()?;
273 ///
274 /// // Create a table
275 /// session.execute(
276 /// "CREATE TABLE test (id UInt64) ENGINE = MergeTree() ORDER BY id",
277 /// None
278 /// )?;
279 ///
280 /// // Query with JSON output
281 /// let result = session.execute(
282 /// "SELECT * FROM test",
283 /// Some(&[Arg::OutputFormat(OutputFormat::JSONEachRow)])
284 /// )?;
285 /// # Ok::<(), chdb_rust::error::Error>(())
286 /// ```
287 ///
288 /// # Errors
289 ///
290 /// Returns an error if:
291 /// - The query syntax is invalid
292 /// - The query references non-existent tables or columns
293 /// - The query execution fails for any other reason
294 pub fn execute(&self, query: &str, query_args: Option<&[Arg]>) -> Result<QueryResult, Error> {
295 let fmt = query_args
296 .and_then(|args| args.iter().find_map(|a| a.as_output_format()))
297 .unwrap_or(self.default_format);
298 self.conn.query(query, fmt)
299 }
300}
301
302impl Drop for Session {
303 fn drop(&mut self) {
304 if self.auto_cleanup {
305 fs::remove_dir_all(&self.data_path).ok();
306 }
307 }
308}