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};
23
24#[pyclass]
25pub struct RiPyORMRepository {
26    pool: RiDatabasePool,
27    table_name: String,
28    rt: &'static tokio::runtime::Runtime,
29}
30
31#[pymethods]
32impl RiPyORMRepository {
33    #[new]
34    #[pyo3(signature = (pool, table_name))]
35    pub fn new(pool: RiDatabasePool, table_name: &str) -> PyResult<Self> {
36        // Security: Validate table name to prevent SQL injection
37        validate_identifier(table_name)
38            .map_err(|e| pyo3::exceptions::PyValueError::new_err(
39                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)
40            ))?;
41        
42        let rt = crate::py_runtime::py_runtime();
43        Ok(Self {
44            pool,
45            table_name: table_name.to_string(),
46            rt,
47        })
48    }
49
50    pub fn get_table_name(&self) -> &str {
51        &self.table_name
52    }
53
54    pub fn find_all(&self, py: Python) -> PyResult<Py<PyList>> {
55        let result = self.rt.block_on(async {
56            self.find_all_impl().await
57        }).map_err(|e| pyo3::PyErr::from(e))?;
58        let list = PyList::empty(py);
59        for value in result {
60            let dict = PyDict::new(py);
61            if let serde_json::Value::Object(map) = value {
62                for (k, v) in map {
63                    Self::json_to_py(py, k, v, &dict);
64                }
65            }
66            list.append(dict).map_err(|e| pyo3::PyErr::from(e))?;
67        }
68        Ok(list.into())
69    }
70
71    pub fn find_by_id(&self, id: &str, py: Python) -> PyResult<Option<Py<PyDict>>> {
72        let result = self.rt.block_on(async {
73            self.find_by_id_impl(id).await
74        }).map_err(|e| pyo3::PyErr::from(e))?;
75        match result {
76            Some(value) => {
77                let dict = PyDict::new(py);
78                if let serde_json::Value::Object(map) = value {
79                    for (k, v) in map {
80                        Self::json_to_py(py, k, v, &dict);
81                    }
82                }
83                Ok(Some(dict.into()))
84            }
85            None => Ok(None),
86        }
87    }
88
89    pub fn count(&self) -> PyResult<u64> {
90        self.rt.block_on(async {
91            self.count_impl().await
92        }).map_err(|e| pyo3::PyErr::from(e))
93    }
94
95    pub fn exists(&self, id: &str) -> PyResult<bool> {
96        self.rt.block_on(async {
97            self.exists_impl(id).await
98        }).map_err(|e| pyo3::PyErr::from(e))
99    }
100
101    pub fn delete(&self, id: &str) -> PyResult<()> {
102        self.rt.block_on(async {
103            self.delete_impl(id).await
104        }).map_err(|e| pyo3::PyErr::from(e))
105    }
106}
107
108impl RiPyORMRepository {
109    fn json_to_py(py: Python, key: String, value: serde_json::Value, dict: &Bound<PyDict>) {
110        match value {
111            serde_json::Value::Null => {},
112            serde_json::Value::Bool(b) => { let _ = dict.set_item(key, b); },
113            serde_json::Value::Number(n) => {
114                if let Some(i) = n.as_i64() {
115                    let _ = dict.set_item(key, i);
116                } else if let Some(f) = n.as_f64() {
117                    let _ = dict.set_item(key, f);
118                } else {
119                    let _ = dict.set_item(key, n.to_string());
120                }
121            }
122            serde_json::Value::String(s) => { let _ = dict.set_item(key, s); },
123            serde_json::Value::Array(arr) => {
124                let list = PyList::empty(py);
125                for (idx, v) in arr.into_iter().enumerate() {
126                    let item = PyDict::new(py);
127                    Self::json_to_py(py, idx.to_string(), v, &item);
128                    let _ = list.append(item);
129                }
130                let _ = dict.set_item(key, list);
131            }
132            serde_json::Value::Object(map) => {
133                let nested = PyDict::new(py);
134                for (k, v) in map {
135                    Self::json_to_py(py, k, v, &nested);
136                }
137                let _ = dict.set_item(key, nested);
138            }
139        }
140    }
141
142    async fn find_all_impl(&self) -> RiResult<Vec<serde_json::Value>> {
143        let db = self.pool.get().await?;
144        let sql = format!("SELECT * FROM {}", self.table_name);
145        let result = db.query(&sql).await?;
146        let mut entities = Vec::with_capacity(4);
147        for row in result {
148            let json_value = serde_json::to_value(row.to_map())?;
149            entities.push(json_value);
150        }
151        Ok(entities)
152    }
153
154    async fn find_by_id_impl(&self, id: &str) -> RiResult<Option<serde_json::Value>> {
155        let db = self.pool.get().await?;
156        let sql = format!("SELECT * FROM {} WHERE id = ?", self.table_name);
157        let result = db.query_with_params(&sql, &[serde_json::json!(id)]).await?;
158        if let Some(row) = result.first() {
159            let json_value = serde_json::to_value(row.to_map())?;
160            Ok(Some(json_value))
161        } else {
162            Ok(None)
163        }
164    }
165
166    async fn count_impl(&self) -> RiResult<u64> {
167        let db = self.pool.get().await?;
168        let sql = format!("SELECT COUNT(*) as total FROM {}", self.table_name);
169        if let Some(row) = db.query_one(&sql).await? {
170            Ok(row.get::<i64>("total").map(|v| v as u64).unwrap_or(0))
171        } else {
172            Ok(0)
173        }
174    }
175
176    async fn exists_impl(&self, id: &str) -> RiResult<bool> {
177        let db = self.pool.get().await?;
178        let sql = format!("SELECT 1 FROM {} WHERE id = ? LIMIT 1", self.table_name);
179        let result = db.query_with_params(&sql, &[serde_json::json!(id)]).await?;
180        Ok(!result.is_empty())
181    }
182
183    async fn delete_impl(&self, id: &str) -> RiResult<()> {
184        let db = self.pool.get().await?;
185        let sql = format!("DELETE FROM {} WHERE id = ?", self.table_name);
186        db.execute_with_params(&sql, &[serde_json::json!(id)]).await?;
187        Ok(())
188    }
189}