Skip to main content

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}