1use ::gluex_ccdb::{
2 context::CCDBContext,
3 data::{self, Data, Value},
4 database::{DirectoryHandle, TypeTableHandle, CCDB},
5 models::{ColumnMeta, ColumnType, TypeTableMeta},
6 CCDBError,
7};
8use chrono::{DateTime, Utc};
9use gluex_core::{
10 parsers::parse_timestamp, run_periods::RESTVersionSelection, utils::resolve_path,
11 GlueXCoreError, RESTVersion, RunNumber,
12};
13use pyo3::{
14 conversion::IntoPyObject,
15 exceptions::PyRuntimeError,
16 prelude::*,
17 types::{PyFloat, PyInt, PyModule, PyString},
18};
19use std::{collections::BTreeMap, env, sync::Arc};
20
21fn py_ccdb_error(err: CCDBError) -> PyErr {
22 PyRuntimeError::new_err(err.to_string())
23}
24
25fn resolve_connection_path(path: Option<String>) -> PyResult<String> {
26 let raw_path = match path {
27 Some(value) if !value.is_empty() => value,
28 _ => env::var("CCDB_CONNECTION").map_err(|_| {
29 PyRuntimeError::new_err("CCDB_CONNECTION is not set and no path was provided")
30 })?,
31 };
32 resolve_path(raw_path)
33 .map(|path| path.to_string_lossy().to_string())
34 .map_err(|err| PyRuntimeError::new_err(err.to_string()))
35}
36
37#[pyclass(name = "ColumnType", module = "gluex_ccdb", skip_from_py_object)]
44#[derive(Clone)]
45pub struct PyColumnType {
46 kind: ColumnType,
47}
48
49#[pymethods]
50impl PyColumnType {
51 #[getter]
53 pub fn name(&self) -> &'static str {
54 self.kind.as_str()
55 }
56 fn __repr__(&self) -> String {
57 format!("ColumnType('{}')", self.kind.as_str())
58 }
59}
60
61impl From<ColumnType> for PyColumnType {
62 fn from(kind: ColumnType) -> Self {
63 Self { kind }
64 }
65}
66
67#[allow(missing_docs)]
68#[pyclass(name = "ColumnMeta", module = "gluex_ccdb", skip_from_py_object)]
69#[derive(Clone)]
70pub struct PyColumnMeta {
71 inner: ColumnMeta,
72}
73
74#[pymethods]
75impl PyColumnMeta {
76 #[getter]
77 fn id(&self) -> i64 {
78 self.inner.id()
79 }
80 #[getter]
81 fn name(&self) -> &str {
82 self.inner.name()
83 }
84 #[getter]
85 fn column_type(&self) -> PyColumnType {
86 self.inner.column_type().into()
87 }
88 #[getter]
89 fn order(&self) -> i64 {
90 self.inner.order()
91 }
92 #[getter]
93 fn comment(&self) -> &str {
94 self.inner.comment()
95 }
96
97 fn __repr__(&self) -> String {
98 format!(
99 "ColumnMeta(name='{}', type='{}', order={})",
100 self.inner.name(),
101 self.inner.column_type().as_str(),
102 self.inner.order()
103 )
104 }
105 fn __str__(&self) -> String {
106 self.__repr__()
107 }
108}
109
110#[pyclass(name = "Column", module = "gluex_ccdb", unsendable)]
119pub struct PyColumn {
120 name: String,
121 column_type: ColumnType,
122 column: Arc<data::Column>,
123}
124
125#[pymethods]
126impl PyColumn {
127 #[getter]
129 pub fn name(&self) -> String {
130 self.name.clone()
131 }
132 #[getter]
134 pub fn column_type(&self) -> PyColumnType {
135 PyColumnType::from(self.column_type)
136 }
137
138 pub fn row(&self, py: Python<'_>, row: usize) -> PyResult<Py<PyAny>> {
155 if row >= self.column.len() {
156 return Err(PyRuntimeError::new_err("row index out of range"));
157 }
158 value_to_py(py, self.column.row(row))
159 }
160
161 pub fn values(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
168 let vals: Vec<Py<PyAny>> = match self.column.as_ref() {
169 data::Column::Int(v) => v
170 .iter()
171 .map(|x| PyInt::new(py, *x).unbind().into())
172 .collect(),
173 data::Column::UInt(v) => v
174 .iter()
175 .map(|x| PyInt::new(py, *x).unbind().into())
176 .collect(),
177 data::Column::Long(v) => v
178 .iter()
179 .map(|x| PyInt::new(py, *x).unbind().into())
180 .collect(),
181 data::Column::ULong(v) => v
182 .iter()
183 .map(|x| PyInt::new(py, *x).unbind().into())
184 .collect(),
185 data::Column::Double(v) => v
186 .iter()
187 .map(|x| PyFloat::new(py, *x).unbind().into())
188 .collect(),
189 data::Column::Bool(v) => v
190 .iter()
191 .map(|x| {
192 let obj = (*x).into_pyobject(py).unwrap();
193 <pyo3::Bound<'_, _> as Clone>::clone(&obj)
194 .into_any()
195 .unbind()
196 })
197 .collect(),
198 data::Column::String(v) => v
199 .iter()
200 .map(|s| PyString::new(py, s).unbind().into())
201 .collect(),
202 };
203 Ok(vals)
204 }
205
206 fn __repr__(&self) -> String {
207 format!(
208 "Column(name='{}', type='{}')",
209 self.name(),
210 self.column_type().name()
211 )
212 }
213 fn __str__(&self) -> String {
214 self.__repr__()
215 }
216}
217
218#[allow(missing_docs)]
219#[pyclass(name = "TypeTableMeta", module = "gluex_ccdb", skip_from_py_object)]
220#[derive(Clone)]
221pub struct PyTypeTableMeta {
222 inner: TypeTableMeta,
223}
224
225#[pymethods]
226impl PyTypeTableMeta {
227 #[getter]
228 fn id(&self) -> i64 {
229 self.inner.id()
230 }
231 #[getter]
232 fn name(&self) -> &str {
233 self.inner.name()
234 }
235 #[getter]
236 fn n_rows(&self) -> i64 {
237 self.inner.n_rows()
238 }
239 #[getter]
240 fn n_columns(&self) -> i64 {
241 self.inner.n_columns()
242 }
243 #[getter]
244 fn comment(&self) -> &str {
245 self.inner.comment()
246 }
247
248 fn __repr__(&self) -> String {
249 format!(
250 "TypeTableMeta(name='{}', id={})",
251 self.inner.name(),
252 self.inner.id()
253 )
254 }
255}
256
257#[pyclass(name = "Data", module = "gluex_ccdb", unsendable)]
270pub struct PyData {
271 inner: Arc<Data>,
272}
273
274#[pymethods]
275impl PyData {
276 #[getter]
278 pub fn n_rows(&self) -> usize {
279 self.inner.n_rows()
280 }
281 #[getter]
283 pub fn n_columns(&self) -> usize {
284 self.inner.n_columns()
285 }
286 #[getter]
288 pub fn column_names(&self) -> Vec<String> {
289 self.inner.column_names().to_vec()
290 }
291 #[getter]
293 pub fn column_types(&self) -> Vec<PyColumnType> {
294 self.inner
295 .column_types()
296 .iter()
297 .copied()
298 .map(PyColumnType::from)
299 .collect()
300 }
301
302 pub fn column(&self, column: Bound<'_, PyAny>) -> PyResult<PyColumn> {
319 let idx = parse_column_index(&self.inner, column)?;
320 let name = self.inner.column_names()[idx].clone();
321 let column = self
322 .inner
323 .column_clone(idx)
324 .ok_or_else(|| PyRuntimeError::new_err("column index out of range"))?;
325 let column_type = self.inner.column_types()[idx];
326 Ok(PyColumn {
327 name,
328 column_type,
329 column: Arc::new(column),
330 })
331 }
332
333 pub fn row(&self, row: usize) -> PyResult<PyRowView> {
350 self.inner.row(row).map_err(py_ccdb_error)?;
351 Ok(PyRowView {
352 data: Arc::clone(&self.inner),
353 row,
354 })
355 }
356
357 pub fn rows(&self) -> PyResult<Vec<PyRowView>> {
364 let n_rows = self.inner.n_rows();
365 let data = Arc::clone(&self.inner);
366 Ok((0..n_rows)
367 .map(|row| PyRowView {
368 data: Arc::clone(&data),
369 row,
370 })
371 .collect())
372 }
373
374 pub fn value(
388 &self,
389 py: Python<'_>,
390 column: Bound<'_, PyAny>,
391 row: usize,
392 ) -> PyResult<Py<PyAny>> {
393 let col_idx = parse_column_index(&self.inner, column)?;
394 match self.inner.value(col_idx, row) {
395 Some(v) => value_to_py(py, v),
396 None => Ok(py.None()),
397 }
398 }
399
400 fn __repr__(&self) -> String {
401 let cols: Vec<String> = self
402 .inner
403 .column_names()
404 .iter()
405 .zip(self.inner.column_types())
406 .map(|(n, t)| format!("{}:{}", n, t.as_str()))
407 .collect();
408 format!(
409 "Data(n_rows={}, n_columns={}, columns=[{}])",
410 self.inner.n_rows(),
411 self.inner.n_columns(),
412 cols.join(", ")
413 )
414 }
415}
416
417#[pyclass(name = "RowView", module = "gluex_ccdb")]
426pub struct PyRowView {
427 data: Arc<Data>,
428 row: usize,
429}
430
431#[pymethods]
432impl PyRowView {
433 #[getter]
435 pub fn n_columns(&self, _py: Python<'_>) -> usize {
436 self.data.n_columns()
437 }
438
439 #[getter]
441 pub fn column_types(&self, _py: Python<'_>) -> Vec<PyColumnType> {
442 self.data
443 .column_types()
444 .iter()
445 .copied()
446 .map(PyColumnType::from)
447 .collect()
448 }
449
450 pub fn value(&self, py: Python<'_>, column: Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
462 let idx = parse_column_index(&self.data, column)?;
463 match self.data.value(idx, self.row) {
464 Some(v) => value_to_py(py, v),
465 None => Ok(py.None()),
466 }
467 }
468
469 pub fn columns(&self, py: Python<'_>) -> PyResult<Vec<(String, PyColumnType, Py<PyAny>)>> {
476 let row = self.data.row(self.row).map_err(py_ccdb_error)?;
477 row.iter_columns()
478 .map(|(name, ty, v)| {
479 Ok((
480 name.to_string(),
481 PyColumnType::from(ty),
482 value_to_py(py, v)?,
483 ))
484 })
485 .collect()
486 }
487
488 fn __repr__(&self) -> String {
489 let cols: Vec<String> = self
490 .data
491 .column_names()
492 .iter()
493 .zip(self.data.column_types())
494 .map(|(n, t)| format!("{}:{}", n, t.as_str()))
495 .collect();
496 format!("RowView(row={}, columns=[{}])", self.row, cols.join(", "))
497 }
498}
499
500#[pyclass(name = "TypeTableHandle", module = "gluex_ccdb", unsendable)]
511pub struct PyTypeTableHandle {
512 inner: TypeTableHandle,
513}
514
515#[pymethods]
516impl PyTypeTableHandle {
517 #[getter]
519 pub fn name(&self) -> &str {
520 self.inner.name()
521 }
522 #[getter]
524 pub fn id(&self) -> i64 {
525 self.inner.id()
526 }
527 #[getter]
529 pub fn meta(&self) -> PyTypeTableMeta {
530 PyTypeTableMeta {
531 inner: self.inner.meta().clone(),
532 }
533 }
534 pub fn full_path(&self) -> String {
536 self.inner.full_path()
537 }
538 pub fn columns(&self) -> PyResult<Vec<PyColumnMeta>> {
545 Ok(self
546 .inner
547 .columns()
548 .map_err(py_ccdb_error)?
549 .into_iter()
550 .map(|m| PyColumnMeta { inner: m })
551 .collect())
552 }
553 #[pyo3(signature = (*, runs=None, variation=None, timestamp=None))]
569 pub fn fetch(
570 &self,
571 runs: Option<Vec<RunNumber>>,
572 variation: Option<String>,
573 timestamp: Option<Bound<'_, PyAny>>,
574 ) -> PyResult<BTreeMap<RunNumber, PyData>> {
575 let ctx = build_context(runs, variation, timestamp)?;
576 Ok(self
577 .inner
578 .fetch(&ctx)
579 .map_err(py_ccdb_error)?
580 .into_iter()
581 .map(|(run, data)| {
582 (
583 run,
584 PyData {
585 inner: Arc::new(data),
586 },
587 )
588 })
589 .collect())
590 }
591
592 #[pyo3(signature = (*, run_period, rest_version=None, variation=None, timestamp=None))]
610 pub fn fetch_run_period(
611 &self,
612 run_period: &str,
613 rest_version: Option<Bound<'_, PyAny>>,
614 variation: Option<String>,
615 timestamp: Option<Bound<'_, PyAny>>,
616 ) -> PyResult<BTreeMap<RunNumber, PyData>> {
617 let run_period = run_period
618 .parse()
619 .map_err(|e: GlueXCoreError| py_ccdb_error(CCDBError::GlueXCoreError(e)))?;
620 let rest_version = parse_py_rest_version_selection(run_period, rest_version)?;
621 let mut ctx = CCDBContext::default()
622 .with_run_period(run_period, rest_version)
623 .map_err(py_ccdb_error)?;
624 if let Some(variation) = variation {
625 ctx.variation = variation;
626 }
627 if let Some(ts) = parse_py_timestamp(timestamp)? {
628 ctx.timestamp = ts;
629 }
630 Ok(self
631 .inner
632 .fetch(&ctx)
633 .map_err(py_ccdb_error)?
634 .into_iter()
635 .map(|(run, data)| {
636 (
637 run,
638 PyData {
639 inner: Arc::new(data),
640 },
641 )
642 })
643 .collect())
644 }
645
646 fn __repr__(&self) -> String {
647 format!("TypeTable(\"{}\")", self.inner.full_path())
648 }
649 fn __str__(&self) -> String {
650 self.__repr__()
651 }
652}
653
654#[pyclass(name = "DirectoryHandle", module = "gluex_ccdb", unsendable)]
661pub struct PyDirectoryHandle {
662 inner: DirectoryHandle,
663}
664
665#[pymethods]
666impl PyDirectoryHandle {
667 pub fn full_path(&self) -> String {
669 self.inner.full_path()
670 }
671 pub fn parent(&self) -> Option<Self> {
678 self.inner.parent().map(|inner| Self { inner })
679 }
680 pub fn dirs(&self) -> Vec<Self> {
687 self.inner
688 .dirs()
689 .into_iter()
690 .map(|inner| Self { inner })
691 .collect()
692 }
693 pub fn dir(&self, name: &str) -> PyResult<Self> {
705 Ok(Self {
706 inner: self.inner.dir(name).map_err(py_ccdb_error)?,
707 })
708 }
709 pub fn tables(&self) -> Vec<PyTypeTableHandle> {
716 self.inner
717 .tables()
718 .into_iter()
719 .map(|inner| PyTypeTableHandle { inner })
720 .collect()
721 }
722 pub fn table(&self, name: &str) -> PyResult<PyTypeTableHandle> {
734 Ok(PyTypeTableHandle {
735 inner: self.inner.table(name).map_err(py_ccdb_error)?,
736 })
737 }
738 fn __repr__(&self) -> String {
739 format!("Directory(\"{}\")", self.full_path())
740 }
741 fn __str__(&self) -> String {
742 self.__repr__()
743 }
744}
745
746#[pyclass(name = "CCDB", module = "gluex_ccdb", unsendable)]
754pub struct PyCCDB {
755 inner: CCDB,
756}
757
758#[pymethods]
759impl PyCCDB {
760 #[new]
768 #[pyo3(signature = (path=None))]
769 pub fn new(path: Option<String>) -> PyResult<Self> {
770 let path = resolve_connection_path(path)?;
771 Ok(Self {
772 inner: CCDB::open(path).map_err(py_ccdb_error)?,
773 })
774 }
775
776 pub fn dir(&self, path: &str) -> PyResult<PyDirectoryHandle> {
788 Ok(PyDirectoryHandle {
789 inner: self.inner.dir(path).map_err(py_ccdb_error)?,
790 })
791 }
792 pub fn table(&self, path: &str) -> PyResult<PyTypeTableHandle> {
804 Ok(PyTypeTableHandle {
805 inner: self.inner.table(path).map_err(py_ccdb_error)?,
806 })
807 }
808 #[pyo3(signature = (path, *, runs=None, variation=None, timestamp=None))]
826 pub fn fetch(
827 &self,
828 path: &str,
829 runs: Option<Vec<RunNumber>>,
830 variation: Option<String>,
831 timestamp: Option<Bound<'_, PyAny>>,
832 ) -> PyResult<BTreeMap<RunNumber, PyData>> {
833 let ctx = build_context(runs, variation, timestamp)?;
834 Ok(self
835 .inner
836 .fetch(path, &ctx)
837 .map_err(py_ccdb_error)?
838 .into_iter()
839 .map(|(run, data)| {
840 (
841 run,
842 PyData {
843 inner: Arc::new(data),
844 },
845 )
846 })
847 .collect())
848 }
849
850 #[pyo3(signature = (path, *, run_period, rest_version=None, variation=None, timestamp=None))]
870 pub fn fetch_run_period(
871 &self,
872 path: &str,
873 run_period: &str,
874 rest_version: Option<Bound<'_, PyAny>>,
875 variation: Option<String>,
876 timestamp: Option<Bound<'_, PyAny>>,
877 ) -> PyResult<BTreeMap<RunNumber, PyData>> {
878 let run_period = run_period
879 .parse()
880 .map_err(|e: GlueXCoreError| py_ccdb_error(CCDBError::GlueXCoreError(e)))?;
881 let rest_version = parse_py_rest_version_selection(run_period, rest_version)?;
882 let mut ctx = CCDBContext::default()
883 .with_run_period(run_period, rest_version)
884 .map_err(py_ccdb_error)?;
885 if let Some(variation) = variation {
886 ctx.variation = variation;
887 }
888 if let Some(ts) = parse_py_timestamp(timestamp)? {
889 ctx.timestamp = ts;
890 }
891 Ok(self
892 .inner
893 .fetch(path, &ctx)
894 .map_err(py_ccdb_error)?
895 .into_iter()
896 .map(|(run, data)| {
897 (
898 run,
899 PyData {
900 inner: Arc::new(data),
901 },
902 )
903 })
904 .collect())
905 }
906
907 pub fn root(&self) -> PyResult<PyDirectoryHandle> {
914 Ok(PyDirectoryHandle {
915 inner: self.inner.root(),
916 })
917 }
918 #[getter]
920 pub fn connection_path(&self) -> &str {
921 self.inner.connection_path()
922 }
923
924 fn __repr__(&self) -> String {
925 format!("CCDB(\"{}\")", self.inner.connection_path())
926 }
927 fn __str__(&self) -> String {
928 self.__repr__()
929 }
930}
931
932fn value_to_py(py: Python<'_>, value: Value<'_>) -> PyResult<Py<PyAny>> {
933 Ok(match value {
934 Value::Int(v) => PyInt::new(py, *v).unbind().into(),
935 Value::UInt(v) => PyInt::new(py, *v).unbind().into(),
936 Value::Long(v) => PyInt::new(py, *v).unbind().into(),
937 Value::ULong(v) => PyInt::new(py, *v).unbind().into(),
938 Value::Double(v) => PyFloat::new(py, *v).unbind().into(),
939 Value::Bool(v) => {
940 let obj = (*v).into_pyobject(py)?;
941 <pyo3::Bound<'_, _> as Clone>::clone(&obj)
942 .into_any()
943 .unbind()
944 }
945 Value::String(v) => PyString::new(py, v).unbind().into(),
946 })
947}
948
949fn parse_py_timestamp(ts: Option<Bound<'_, PyAny>>) -> PyResult<Option<DateTime<Utc>>> {
950 let Some(val) = ts else {
951 return Ok(None);
952 };
953 if let Ok(dt) = val.extract::<DateTime<Utc>>() {
954 return Ok(Some(dt));
955 }
956 if let Ok(s) = val.extract::<String>() {
957 let parsed = parse_timestamp(&s).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
958 return Ok(Some(parsed));
959 }
960 Err(PyRuntimeError::new_err("timestamp must be str or datetime"))
961}
962
963fn parse_py_rest_version_selection(
964 run_period: gluex_core::run_periods::RunPeriod,
965 rest_version: Option<Bound<'_, PyAny>>,
966) -> PyResult<RESTVersionSelection> {
967 let Some(val) = rest_version else {
968 return Ok(RESTVersionSelection::Current);
969 };
970 if let Ok(version) = val.extract::<RESTVersion>() {
971 return RESTVersionSelection::try_new(run_period, version)
972 .map_err(|e| PyRuntimeError::new_err(e.to_string()));
973 }
974 if let Ok(timestamp) = val.extract::<DateTime<Utc>>() {
975 return Ok(RESTVersionSelection::from_timestamp(timestamp));
976 }
977 Err(PyRuntimeError::new_err(
978 "rest_version must be int, datetime, or None",
979 ))
980}
981
982fn parse_column_index(data: &Data, column: Bound<'_, PyAny>) -> PyResult<usize> {
983 if let Ok(idx) = column.extract::<usize>() {
984 if idx < data.n_columns() {
985 return Ok(idx);
986 }
987 return Err(PyRuntimeError::new_err("column index out of range"));
988 }
989 if let Ok(name) = column.extract::<String>() {
990 if let Some(idx) = data.column_names().iter().position(|n| n == &name) {
991 return Ok(idx);
992 }
993 return Err(PyRuntimeError::new_err("column name not found"));
994 }
995 Err(PyRuntimeError::new_err("column must be int or str"))
996}
997
998fn build_context(
999 runs: Option<Vec<RunNumber>>,
1000 variation: Option<String>,
1001 timestamp: Option<Bound<'_, PyAny>>,
1002) -> PyResult<CCDBContext> {
1003 let mut ctx = CCDBContext::default();
1004 if let Some(runs) = runs {
1005 ctx.runs = runs;
1006 }
1007 if let Some(variation) = variation {
1008 ctx.variation = variation;
1009 }
1010 if let Some(ts) = parse_py_timestamp(timestamp)? {
1011 ctx.timestamp = ts;
1012 }
1013 Ok(ctx)
1014}
1015
1016#[pymodule]
1017pub fn gluex_ccdb(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
1019 m.add_class::<PyCCDB>()?;
1020 m.add_class::<PyTypeTableHandle>()?;
1021 m.add_class::<PyDirectoryHandle>()?;
1022 m.add_class::<PyData>()?;
1023 m.add_class::<PyRowView>()?;
1024 m.add_class::<PyColumn>()?;
1025 m.add_class::<PyColumnMeta>()?;
1026 m.add_class::<PyTypeTableMeta>()?;
1027 m.add_class::<PyColumnType>()?;
1028 Ok(())
1029}