Skip to main content

ri/database/orm/
py_repository.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18use crate::core::RiResult;
19use crate::database::{RiDatabase, RiDatabasePool};
20use crate::database::orm::validate_identifier;
21use pyo3::prelude::*;
22use pyo3::types::{PyDict, PyList};
23use tokio::runtime::Runtime;
24
25#[pyclass]
26pub struct RiPyORMRepository {
27    pool: RiDatabasePool,
28    table_name: String,
29    rt: Runtime,
30}
31
32#[pymethods]
33impl RiPyORMRepository {
34    #[new]
35    #[pyo3(signature = (pool, table_name))]
36    pub fn new(pool: RiDatabasePool, table_name: &str) -> PyResult<Self> {
37        // Security: Validate table name to prevent SQL injection
38        validate_identifier(table_name)
39            .map_err(|e| pyo3::exceptions::PyValueError::new_err(
40                format!("Invalid table name '{}': {}. Table names must contain only alphanumeric characters and underscores, start with a letter or underscore, and be 1-128 characters long.", table_name, e)
41            ))?;
42        
43        let rt = Runtime::new().map_err(|e| pyo3::PyErr::from(crate::core::RiError::Other(e.to_string())))?;
44        Ok(Self {
45            pool,
46            table_name: table_name.to_string(),
47            rt,
48        })
49    }
50
51    pub fn get_table_name(&self) -> &str {
52        &self.table_name
53    }
54
55    pub fn find_all(&self, py: Python) -> PyResult<Py<PyList>> {
56        let result = self.rt.block_on(async {
57            self.find_all_impl().await
58        }).map_err(|e| pyo3::PyErr::from(e))?;
59        let list = PyList::empty(py);
60        for value in result {
61            let dict = PyDict::new(py);
62            if let serde_json::Value::Object(map) = value {
63                for (k, v) in map {
64                    Self::json_to_py(py, k, v, &dict);
65                }
66            }
67            list.append(dict).map_err(|e| pyo3::PyErr::from(e))?;
68        }
69        Ok(list.into())
70    }
71
72    pub fn find_by_id(&self, id: &str, py: Python) -> PyResult<Option<Py<PyDict>>> {
73        let result = self.rt.block_on(async {
74            self.find_by_id_impl(id).await
75        }).map_err(|e| pyo3::PyErr::from(e))?;
76        match result {
77            Some(value) => {
78                let dict = PyDict::new(py);
79                if let serde_json::Value::Object(map) = value {
80                    for (k, v) in map {
81                        Self::json_to_py(py, k, v, &dict);
82                    }
83                }
84                Ok(Some(dict.into()))
85            }
86            None => Ok(None),
87        }
88    }
89
90    pub fn count(&self) -> PyResult<u64> {
91        self.rt.block_on(async {
92            self.count_impl().await
93        }).map_err(|e| pyo3::PyErr::from(e))
94    }
95
96    pub fn exists(&self, id: &str) -> PyResult<bool> {
97        self.rt.block_on(async {
98            self.exists_impl(id).await
99        }).map_err(|e| pyo3::PyErr::from(e))
100    }
101
102    pub fn delete(&self, id: &str) -> PyResult<()> {
103        self.rt.block_on(async {
104            self.delete_impl(id).await
105        }).map_err(|e| pyo3::PyErr::from(e))
106    }
107}
108
109impl RiPyORMRepository {
110    fn json_to_py(py: Python, key: String, value: serde_json::Value, dict: &Bound<PyDict>) {
111        match value {
112            serde_json::Value::Null => {},
113            serde_json::Value::Bool(b) => { let _ = dict.set_item(key, b); },
114            serde_json::Value::Number(n) => {
115                if let Some(i) = n.as_i64() {
116                    let _ = dict.set_item(key, i);
117                } else if let Some(f) = n.as_f64() {
118                    let _ = dict.set_item(key, f);
119                } else {
120                    let _ = dict.set_item(key, n.to_string());
121                }
122            }
123            serde_json::Value::String(s) => { let _ = dict.set_item(key, s); },
124            serde_json::Value::Array(arr) => {
125                let list = PyList::empty(py);
126                for (idx, v) in arr.into_iter().enumerate() {
127                    let item = PyDict::new(py);
128                    Self::json_to_py(py, idx.to_string(), v, &item);
129                    let _ = list.append(item);
130                }
131                let _ = dict.set_item(key, list);
132            }
133            serde_json::Value::Object(map) => {
134                let nested = PyDict::new(py);
135                for (k, v) in map {
136                    Self::json_to_py(py, k, v, &nested);
137                }
138                let _ = dict.set_item(key, nested);
139            }
140        }
141    }
142
143    async fn find_all_impl(&self) -> RiResult<Vec<serde_json::Value>> {
144        let db = self.pool.get().await?;
145        let sql = format!("SELECT * FROM {}", self.table_name);
146        let result = db.query(&sql).await?;
147        let mut entities = Vec::with_capacity(4);
148        for row in result {
149            let json_value = serde_json::to_value(row.to_map())?;
150            entities.push(json_value);
151        }
152        Ok(entities)
153    }
154
155    async fn find_by_id_impl(&self, id: &str) -> RiResult<Option<serde_json::Value>> {
156        let db = self.pool.get().await?;
157        let sql = format!("SELECT * FROM {} WHERE id = ?", self.table_name);
158        let result = db.query_with_params(&sql, &[serde_json::json!(id)]).await?;
159        if let Some(row) = result.first() {
160            let json_value = serde_json::to_value(row.to_map())?;
161            Ok(Some(json_value))
162        } else {
163            Ok(None)
164        }
165    }
166
167    async fn count_impl(&self) -> RiResult<u64> {
168        let db = self.pool.get().await?;
169        let sql = format!("SELECT COUNT(*) as total FROM {}", self.table_name);
170        if let Some(row) = db.query_one(&sql).await? {
171            Ok(row.get::<i64>("total").map(|v| v as u64).unwrap_or(0))
172        } else {
173            Ok(0)
174        }
175    }
176
177    async fn exists_impl(&self, id: &str) -> RiResult<bool> {
178        let db = self.pool.get().await?;
179        let sql = format!("SELECT 1 FROM {} WHERE id = ? LIMIT 1", self.table_name);
180        let result = db.query_with_params(&sql, &[serde_json::json!(id)]).await?;
181        Ok(!result.is_empty())
182    }
183
184    async fn delete_impl(&self, id: &str) -> RiResult<()> {
185        let db = self.pool.get().await?;
186        let sql = format!("DELETE FROM {} WHERE id = ?", self.table_name);
187        db.execute_with_params(&sql, &[serde_json::json!(id)]).await?;
188        Ok(())
189    }
190}