polars_python/
sql.rs

1use parking_lot::RwLock;
2use polars::sql::SQLContext;
3use pyo3::prelude::*;
4
5use crate::PyLazyFrame;
6use crate::error::PyPolarsErr;
7
8#[pyclass(frozen)]
9#[repr(transparent)]
10pub struct PySQLContext {
11    pub context: RwLock<SQLContext>,
12}
13
14impl Clone for PySQLContext {
15    fn clone(&self) -> Self {
16        Self {
17            context: RwLock::new(self.context.read().clone()),
18        }
19    }
20}
21
22#[pymethods]
23#[allow(
24    clippy::wrong_self_convention,
25    clippy::should_implement_trait,
26    clippy::len_without_is_empty
27)]
28impl PySQLContext {
29    #[staticmethod]
30    #[allow(clippy::new_without_default)]
31    pub fn new() -> PySQLContext {
32        PySQLContext {
33            context: RwLock::new(SQLContext::new()),
34        }
35    }
36
37    pub fn execute(&self, query: &str) -> PyResult<PyLazyFrame> {
38        Ok(self
39            .context
40            .write()
41            .execute(query)
42            .map_err(PyPolarsErr::from)?
43            .into())
44    }
45
46    pub fn get_tables(&self) -> PyResult<Vec<String>> {
47        Ok(self.context.read().get_tables())
48    }
49
50    pub fn register(&self, name: &str, lf: PyLazyFrame) {
51        self.context.write().register(name, lf.ldf.into_inner())
52    }
53
54    pub fn unregister(&self, name: &str) {
55        self.context.write().unregister(name)
56    }
57}