polars_python/
sql.rs

1use parking_lot::RwLock;
2use polars::sql::{SQLContext, extract_table_identifiers};
3use pyo3::prelude::*;
4
5use crate::PyLazyFrame;
6use crate::error::PyPolarsErr;
7use crate::utils::EnterPolarsExt;
8
9#[pyclass(frozen, skip_from_py_object)]
10#[repr(transparent)]
11pub struct PySQLContext {
12    pub context: RwLock<SQLContext>,
13}
14
15impl Clone for PySQLContext {
16    fn clone(&self) -> Self {
17        Self {
18            context: RwLock::new(self.context.read().clone()),
19        }
20    }
21}
22
23#[pymethods]
24#[allow(
25    clippy::wrong_self_convention,
26    clippy::should_implement_trait,
27    clippy::len_without_is_empty
28)]
29impl PySQLContext {
30    #[staticmethod]
31    #[allow(clippy::new_without_default)]
32    pub fn new() -> PySQLContext {
33        PySQLContext {
34            context: RwLock::new(SQLContext::new()),
35        }
36    }
37
38    /// Execute a SQL query in the current SQLContext.
39    pub fn execute(&self, py: Python<'_>, query: &str) -> PyResult<PyLazyFrame> {
40        py.enter_polars(|| self.context.write().execute(query))
41            .map(Into::into)
42    }
43
44    /// Get a list of table names registered in the current SQLContext.
45    pub fn get_tables(&self) -> PyResult<Vec<String>> {
46        Ok(self.context.read().get_tables())
47    }
48
49    /// Register a table in the current SQLContext.
50    pub fn register(&self, name: &str, lf: PyLazyFrame) {
51        self.context.write().register(name, lf.ldf.into_inner())
52    }
53
54    /// Unregister a table from the current SQLContext.
55    pub fn unregister(&self, name: &str) {
56        self.context.write().unregister(name)
57    }
58
59    /// Extract table identifiers from a SQL query string.
60    #[staticmethod]
61    #[pyo3(signature = (query, include_schema=true, unique=false))]
62    pub fn table_identifiers(
63        query: &str,
64        include_schema: bool,
65        unique: bool,
66    ) -> PyResult<Vec<String>> {
67        extract_table_identifiers(query, include_schema, unique)
68            .map_err(PyPolarsErr::from)
69            .map_err(Into::into)
70    }
71}
polars_python/sql.rs

polars_python/
sql.rs
