polars_python/
sql.rs

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