graphrecords_python/graphrecord/
connector.rs

1use crate::prelude::PyGraphRecord;
2use graphrecords_core::{
3    GraphRecord,
4    errors::{GraphRecordError, GraphRecordResult},
5    graphrecord::connector::{Connector, ExportConnector, IngestConnector},
6};
7use pyo3::{Py, PyAny, Python, types::PyAnyMethods};
8use serde::{Deserialize, Deserializer, Serialize, Serializer};
9
10#[derive(Debug)]
11pub struct PyConnector(Py<PyAny>);
12
13impl PyConnector {
14    pub const fn new(connector: Py<PyAny>) -> Self {
15        Self(connector)
16    }
17}
18
19impl Clone for PyConnector {
20    fn clone(&self) -> Self {
21        Python::attach(|py| Self(self.0.clone_ref(py)))
22    }
23}
24
25impl Serialize for PyConnector {
26    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
27        Python::attach(|py| {
28            let cloudpickle = py
29                .import("cloudpickle")
30                .map_err(serde::ser::Error::custom)?;
31
32            let bytes: Vec<u8> = cloudpickle
33                .call_method1("dumps", (&self.0,))
34                .map_err(serde::ser::Error::custom)?
35                .extract()
36                .map_err(serde::ser::Error::custom)?;
37
38            serializer.serialize_bytes(&bytes)
39        })
40    }
41}
42
43impl<'de> Deserialize<'de> for PyConnector {
44    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
45        let bytes: Vec<u8> = Deserialize::deserialize(deserializer)?;
46
47        Python::attach(|py| {
48            let cloudpickle = py.import("cloudpickle").map_err(serde::de::Error::custom)?;
49
50            let obj: Py<PyAny> = cloudpickle
51                .call_method1("loads", (bytes.as_slice(),))
52                .map_err(serde::de::Error::custom)?
53                .into();
54
55            Ok(Self(obj))
56        })
57    }
58}
59
60impl Connector for PyConnector {
61    fn initialize(&self, graphrecord: &mut GraphRecord) -> GraphRecordResult<()> {
62        Python::attach(|py| {
63            PyGraphRecord::scope_mut(py, graphrecord, |py, graphrecord| {
64                self.0
65                    .call_method1(py, "initialize", (graphrecord,))
66                    .map_err(|err| GraphRecordError::ConnectorFailure {
67                        message: err.to_string(),
68                    })?;
69
70                Ok(())
71            })
72        })
73    }
74
75    fn disconnect(&self, graphrecord: &mut GraphRecord) -> GraphRecordResult<()> {
76        Python::attach(|py| {
77            PyGraphRecord::scope_mut(py, graphrecord, |py, graphrecord| {
78                self.0
79                    .call_method1(py, "disconnect", (graphrecord,))
80                    .map_err(|err| GraphRecordError::ConnectorFailure {
81                        message: err.to_string(),
82                    })?;
83
84                Ok(())
85            })
86        })
87    }
88}
89
90impl IngestConnector for PyConnector {
91    type DataSet = Py<PyAny>;
92
93    fn ingest(&self, graphrecord: &mut GraphRecord, data: Self::DataSet) -> GraphRecordResult<()> {
94        Python::attach(|py| {
95            PyGraphRecord::scope_mut(py, graphrecord, |py, graphrecord| {
96                self.0
97                    .call_method1(py, "ingest", (graphrecord, data))
98                    .map_err(|err| GraphRecordError::ConnectorFailure {
99                        message: err.to_string(),
100                    })?;
101
102                Ok(())
103            })
104        })
105    }
106}
107
108impl ExportConnector for PyConnector {
109    type DataSet = Py<PyAny>;
110
111    fn export(&self, graphrecord: &GraphRecord) -> GraphRecordResult<Self::DataSet> {
112        Python::attach(|py| {
113            PyGraphRecord::scope(py, graphrecord, |py, graphrecord| {
114                let data = self
115                    .0
116                    .call_method1(py, "export", (graphrecord,))
117                    .map_err(|err| GraphRecordError::ConnectorFailure {
118                        message: err.to_string(),
119                    })?;
120
121                Ok(data)
122            })
123        })
124    }
125}
graphrecords_python/graphrecord/connector.rs

graphrecords_python/graphrecord/
connector.rs
