fluidB
/
fluidb
1
0
Fork 0
Code Pull Requests Activity

improve python API

Browse Source
main
Ziyang Hu 2 years ago
parent 00825ac39e
commit 6411e42e9c
3 changed files with 81 additions and 42 deletions
Show Stats Download Patch File Download Diff File

14
cozo-core/src/lib.rs
Unescape Escape View File

@ -222,10 +222,11 @@ impl DbInstance {
self.run_script_fold_err(payload, params_json).to_string()
}
/// Dispatcher method. See [crate::Db::export_relations].
pub fn export_relations<'a>(
&self,
relations: impl Iterator<Item = &'a str>,
) -> Result<BTreeMap<String, NamedRows>> {
pub fn export_relations<'a, I, T>(&self, relations: I) -> Result<BTreeMap<String, NamedRows>>
where
T: AsRef<str>,
I: Iterator<Item = T>,
{
match self {
DbInstance::Mem(db) => db.export_relations(relations),
#[cfg(feature = "storage-sqlite")]
@ -413,7 +414,10 @@ impl DbInstance {
}
}
/// Dispatcher method. See [crate::Db::register_fixed_rule].
pub fn register_fixed_rule<R>(&self, name: String, rule_impl: R) -> Result<()> where R: FixedRule + 'static {
pub fn register_fixed_rule<R>(&self, name: String, rule_impl: R) -> Result<()>
where
R: FixedRule + 'static,
{
match self {
DbInstance::Mem(db) => db.register_fixed_rule(name, rule_impl),
#[cfg(feature = "storage-sqlite")]

16
cozo-core/src/runtime/db.rs
Unescape Escape View File

@ -118,7 +118,8 @@ pub struct NamedRows {
pub headers: Vec<String>,
/// The rows
pub rows: Vec<Tuple>,
pub(crate) next: Option<Box<NamedRows>>,
/// Contains the next named rows, if exists
pub next: Option<Box<NamedRows>>,
}
impl NamedRows {
@ -214,14 +215,15 @@ impl<'s, S: Storage<'s>> Db<S> {
/// Export relations to JSON data.
///
/// `relations` contains names of the stored relations to export.
pub fn export_relations<'a>(
&'s self,
relations: impl Iterator<Item = &'a str>,
) -> Result<BTreeMap<String, NamedRows>> {
pub fn export_relations<'a, I, T>(&'s self, relations: I) -> Result<BTreeMap<String, NamedRows>>
where
T: AsRef<str>,
I: Iterator<Item = T>,
{
let tx = self.transact()?;
let mut ret: BTreeMap<String, NamedRows> = BTreeMap::new();
for rel in relations {
let handle = tx.get_relation(rel, false)?;
let handle = tx.get_relation(rel.as_ref(), false)?;
if handle.access_level < AccessLevel::ReadOnly {
bail!(InsufficientAccessLevel(
@ -256,7 +258,7 @@ impl<'s, S: Storage<'s>> Db<S> {
rows.push(tuple);
}
let headers = cols.iter().map(|col| col.to_string()).collect_vec();
ret.insert(rel.to_string(), NamedRows::new(headers, rows));
ret.insert(rel.as_ref().to_string(), NamedRows::new(headers, rows));
}
Ok(ret)
}

91
cozo-lib-python/src/lib.rs
Unescape Escape View File

@ -9,21 +9,37 @@
use std::collections::BTreeMap;
use std::thread;
use miette::{IntoDiagnostic, Result};
use miette::{IntoDiagnostic, Report, Result};
use pyo3::exceptions::PyException;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyString, PyTuple};
use cozo::*;
fn py_to_named_rows(ob: &PyAny) -> PyResult<NamedRows> {
fn py_to_rows(ob: &PyAny) -> PyResult<Vec<Vec<DataValue>>> {
let rows = ob.extract::<Vec<Vec<&PyAny>>>()?;
let res: Vec<Vec<DataValue>> = rows
.into_iter()
.map(|row| row.into_iter().map(py_to_value).collect::<PyResult<_>>())
.collect::<PyResult<_>>()?;
Ok(res)
}
Ok(NamedRows::new(vec![], res))
fn report2py(r: Report) -> PyErr {
PyException::new_err(r.to_string())
}
fn py_to_named_rows(ob: &PyAny) -> PyResult<NamedRows> {
let d = ob.downcast::<PyDict>()?;
let rows = d
.get_item("rows")
.ok_or_else(|| PyException::new_err("named rows must contain 'rows'"))?;
let rows = py_to_rows(rows)?;
let headers = d
.get_item("headers")
.ok_or_else(|| PyException::new_err("named rows must contain 'headers'"))?;
let headers = headers.extract::<Vec<String>>()?;
Ok(NamedRows::new(headers, rows))
}
fn py_to_value(ob: &PyAny) -> PyResult<DataValue> {
@ -76,7 +92,6 @@ fn options_to_py(opts: BTreeMap<String, DataValue>, py: Python<'_>) -> PyResult<
let val = value_to_py(v, py);
ret.set_item(k, val)?;
}
Ok(ret.into())
}
@ -122,7 +137,11 @@ fn rows_to_py_rows(rows: Vec<Vec<DataValue>>, py: Python<'_>) -> PyObject {
fn named_rows_to_py(named_rows: NamedRows, py: Python<'_>) -> PyObject {
let rows = rows_to_py_rows(named_rows.rows, py);
let headers = named_rows.headers.into_py(py);
BTreeMap::from([("rows", rows), ("headers", headers)]).into_py(py)
let next = match named_rows.next {
None => py.None(),
Some(nxt) => named_rows_to_py(*nxt, py),
};
BTreeMap::from([("rows", rows), ("headers", headers), ("next", next)]).into_py(py)
}
#[pyclass]
@ -195,11 +214,10 @@ impl CozoDbPy {
let py_opts = options_to_py(options, py).into_diagnostic()?;
let args = PyTuple::new(py, vec![PyObject::from(py_inputs), py_opts]);
let res = cb.as_ref(py).call1(args).into_diagnostic()?;
py_to_named_rows(res).into_diagnostic()
Ok(NamedRows::new(vec![], py_to_rows(res).into_diagnostic()?))
})
});
db.register_fixed_rule(name, rule_impl)
.map_err(|err| PyException::new_err(err.to_string()))
db.register_fixed_rule(name, rule_impl).map_err(report2py)
} else {
Err(PyException::new_err(DB_CLOSED_MSG))
}
@ -221,46 +239,61 @@ impl CozoDbPy {
Ok(false)
}
}
pub fn run_query(&self, py: Python<'_>, query: &str, params: &str) -> String {
pub fn export_relations(&self, py: Python<'_>, relations: Vec<String>) -> PyResult<PyObject> {
if let Some(db) = &self.db {
py.allow_threads(|| db.run_script_str(query, params))
} else {
DB_CLOSED_MSG.to_string()
}
let res = match py.allow_threads(|| db.export_relations(relations.iter())) {
Ok(res) => res,
Err(err) => return Err(PyException::new_err(err.to_string())),
};
let ret = PyDict::new(py);
for (k, v) in res {
ret.set_item(k, named_rows_to_py(v, py))?;
}
pub fn export_relations(&self, py: Python<'_>, data: &str) -> String {
if let Some(db) = &self.db {
py.allow_threads(|| db.export_relations_str(data))
Ok(ret.into())
} else {
DB_CLOSED_MSG.to_string()
Err(PyException::new_err(DB_CLOSED_MSG.to_string()))
}
}
pub fn import_relations(&self, py: Python<'_>, data: &str) -> String {
pub fn import_relations(&self, py: Python<'_>, data: &PyDict) -> PyResult<()> {
if let Some(db) = &self.db {
py.allow_threads(|| db.import_relations_str(data))
let mut arg = BTreeMap::new();
for (k, v) in data.iter() {
let k = k.extract::<String>()?;
let vals = py_to_named_rows(v)?;
arg.insert(k, vals);
}
py.allow_threads(|| db.import_relations(arg))
.map_err(report2py)
} else {
DB_CLOSED_MSG.to_string()
Err(PyException::new_err(DB_CLOSED_MSG.to_string()))
}
}
pub fn backup(&self, py: Python<'_>, path: &str) -> String {
pub fn backup(&self, py: Python<'_>, path: &str) -> PyResult<()> {
if let Some(db) = &self.db {
py.allow_threads(|| db.backup_db_str(path))
py.allow_threads(|| db.backup_db(path)).map_err(report2py)
} else {
DB_CLOSED_MSG.to_string()
Err(PyException::new_err(DB_CLOSED_MSG.to_string()))
}
}
pub fn restore(&self, py: Python<'_>, path: &str) -> String {
pub fn restore(&self, py: Python<'_>, path: &str) -> PyResult<()> {
if let Some(db) = &self.db {
py.allow_threads(|| db.restore_backup_str(path))
py.allow_threads(|| db.restore_backup(path))
.map_err(report2py)
} else {
DB_CLOSED_MSG.to_string()
Err(PyException::new_err(DB_CLOSED_MSG.to_string()))
}
}
pub fn import_from_backup(&self, py: Python<'_>, data: &str) -> String {
pub fn import_from_backup(
&self,
py: Python<'_>,
in_file: &str,
relations: Vec<String>,
) -> PyResult<()> {
if let Some(db) = &self.db {
py.allow_threads(|| db.import_from_backup_str(data))
py.allow_threads(|| db.import_from_backup(in_file, &relations))
.map_err(report2py)
} else {
DB_CLOSED_MSG.to_string()
Err(PyException::new_err(DB_CLOSED_MSG.to_string()))
}
}
pub fn close(&mut self) -> bool {

Write Preview
Loading…
Cancel
Save