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}