Nov 9, 2023 · Nov 2, 2023 · Nov 2, 2023 · Nov 3, 2023 · Nov 3, 2023 · Nov 6, 2023
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
 use ndarray::Zip;
 use pgrx::iter::{SetOfIterator, TableIterator};
 use pgrx::*;
 use pyo3::prelude::*;
 use pyo3::types::{IntoPyDict, PyDict};

 #[cfg(feature = "python")]
 use serde_json::json;
    }
 }

 struct TransformStreamIterator {
    locals: Py<PyDict>,
 }

 impl TransformStreamIterator {
    fn new(python_iter: Py<PyAny>) -> Self {
        let locals = Python::with_gil(|py| -> Result<Py<PyDict>, PyErr> {
            Ok([("python_iter", python_iter)].into_py_dict(py).into())
        })
        .map_err(|e| error!("{e}"))
        .unwrap();
        Self { locals }
    }
 }

 impl Iterator for TransformStreamIterator {
    type Item = String;
    fn next(&mut self) -> Option<Self::Item> {
        // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
        Python::with_gil(|py| -> Result<Option<String>, PyErr> {
            let code = "next(python_iter)";
            let res: &PyAny = py.eval(code, Some(self.locals.as_ref(py)), None)?;
            if res.is_none() {
                Ok(None)
            } else {
                let res: String = res.extract()?;
                Ok(Some(res))
            }
        })
        .map_err(|e| error!("{e}"))
        .unwrap()
    }
 }

 #[cfg(all(feature = "python", not(feature = "use_as_lib")))]
 #[pg_extern(immutable, parallel_safe, name = "transform_stream")]
 #[allow(unused_variables)] // cache is maintained for api compatibility
 pub fn transform_stream_json(
    task: JsonB,
    args: default!(JsonB, "'{}'"),
    input: default!(&str, "''"),
    cache: default!(bool, false),
 ) -> SetOfIterator<'static, String> {
    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
    let python_iter = crate::bindings::transformers::transform_stream(&task.0, &args.0, input)
        .map_err(|e| error!("{e}"))
        .unwrap();
    let res = TransformStreamIterator::new(python_iter);
    SetOfIterator::new(res)
 }

 #[cfg(all(feature = "python", not(feature = "use_as_lib")))]
 #[pg_extern(immutable, parallel_safe, name = "transform_stream")]
 #[allow(unused_variables)] // cache is maintained for api compatibility
 pub fn transform_stream_string(
    task: String,
    args: default!(JsonB, "'{}'"),
    input: default!(&str, "''"),
    cache: default!(bool, false),
 ) -> SetOfIterator<'static, String> {
    let task_json = json!({ "task": task });
    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
    let python_iter = crate::bindings::transformers::transform_stream(&task_json, &args.0, input)
        .map_err(|e| error!("{e}"))
        .unwrap();
    let res = TransformStreamIterator::new(python_iter);
    SetOfIterator::new(res)
 }

 #[cfg(feature = "python")]
 #[pg_extern(immutable, parallel_safe, name = "generate")]
 fn generate(project_name: &str, inputs: &str, config: default!(JsonB, "'{}'")) -> String {
diff --git a/pgml-extension/src/bindings/transformers/mod.rs b/pgml-extension/src/bindings/transformers/mod.rs

 pub mod whitelist;

 create_pymodule!("/src/bindings/transformers/transformers.py");

 pub fn transform(
    task: &serde_json::Value,
    args: &serde_json::Value,
    inputs: Vec<&str>,
 ) -> Result<serde_json::Value> {
    crate::bindings::python::activate()?;

    whitelist::verify_task(task)?;

    let task = serde_json::to_string(task)?;
    let args = serde_json::to_string(args)?;
    let inputs = serde_json::to_string(&inputs)?;
 mod transformers;
 pub use transformers::*;

    let results = Python::with_gil(|py| -> Result<String> {
        let transform: Py<PyAny> = get_module!(PY_MODULE)
            .getattr(py, "transform")
            .format_traceback(py)?;

        let output = transform
            .call1(
                py,
                PyTuple::new(
                    py,
                    &[task.into_py(py), args.into_py(py), inputs.into_py(py)],
                ),
            )
            .format_traceback(py)?;

        output.extract(py).format_traceback(py)
    })?;

    Ok(serde_json::from_str(&results)?)
 }
 create_pymodule!("/src/bindings/transformers/transformers.py");

 pub fn get_model_from(task: &Value) -> Result<String> {
    Python::with_gil(|py| -> Result<String> {
Original file line number	Diff line number	Diff line change
Expand Up		@@ -4,6 +4,8 @@ use std::str::FromStr;
		use ndarray::Zip;
		use pgrx::iter::{SetOfIterator, TableIterator};
		use pgrx::*;
		use pyo3::prelude::*;
		use pyo3::types::{IntoPyDict, PyDict};

		#[cfg(feature = "python")]
		use serde_json::json;
Expand DownExpand Up		@@ -632,6 +634,75 @@ pub fn transform_string(
		}
		}

		struct TransformStreamIterator {
		locals: Py<PyDict>,
		}

		impl TransformStreamIterator {
		fn new(python_iter: Py<PyAny>) -> Self {
		let locals = Python::with_gil(\|py\| -> Result<Py<PyDict>, PyErr> {
		Ok([("python_iter", python_iter)].into_py_dict(py).into())
		})
		.map_err(\|e\| error!("{e}"))
		.unwrap();
		Self { locals }
		}
		}

		impl Iterator for TransformStreamIterator {
		type Item = String;
		fn next(&mut self) -> Option<Self::Item> {
		// We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
		Python::with_gil(\|py\| -> Result<Option<String>, PyErr> {
		let code = "next(python_iter)";
		let res: &PyAny = py.eval(code, Some(self.locals.as_ref(py)), None)?;
		if res.is_none() {
		Ok(None)
		} else {
		let res: String = res.extract()?;
		Ok(Some(res))
		}
		})
		.map_err(\|e\| error!("{e}"))
		.unwrap()
		}
		}

		#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
		#[pg_extern(immutable, parallel_safe, name = "transform_stream")]
		#[allow(unused_variables)] // cache is maintained for api compatibility
		pub fn transform_stream_json(
		task: JsonB,
		args: default!(JsonB, "'{}'"),
		input: default!(&str, "''"),
		cache: default!(bool, false),
		) -> SetOfIterator<'static, String> {
		// We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
		let python_iter = crate::bindings::transformers::transform_stream(&task.0, &args.0, input)
		.map_err(\|e\| error!("{e}"))
		.unwrap();
		let res = TransformStreamIterator::new(python_iter);
		SetOfIterator::new(res)
		}

		#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
		#[pg_extern(immutable, parallel_safe, name = "transform_stream")]
		#[allow(unused_variables)] // cache is maintained for api compatibility
		pub fn transform_stream_string(
		task: String,
		args: default!(JsonB, "'{}'"),
		input: default!(&str, "''"),
		cache: default!(bool, false),
		) -> SetOfIterator<'static, String> {
		let task_json = json!({ "task": task });
		// We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
		let python_iter = crate::bindings::transformers::transform_stream(&task_json, &args.0, input)
		.map_err(\|e\| error!("{e}"))
		.unwrap();
		let res = TransformStreamIterator::new(python_iter);
		SetOfIterator::new(res)
		}

		#[cfg(feature = "python")]
		#[pg_extern(immutable, parallel_safe, name = "generate")]
		fn generate(project_name: &str, inputs: &str, config: default!(JsonB, "'{}'")) -> String {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -16,41 +16,10 @@ use super::TracebackError;

		pub mod whitelist;

		create_pymodule!("/src/bindings/transformers/transformers.py");

		pub fn transform(
		task: &serde_json::Value,
		args: &serde_json::Value,
		inputs: Vec<&str>,
		) -> Result<serde_json::Value> {
		crate::bindings::python::activate()?;

		whitelist::verify_task(task)?;

		let task = serde_json::to_string(task)?;
		let args = serde_json::to_string(args)?;
		let inputs = serde_json::to_string(&inputs)?;
		mod transformers;
		pub use transformers::*;

		let results = Python::with_gil(\|py\| -> Result<String> {
		let transform: Py<PyAny> = get_module!(PY_MODULE)
		.getattr(py, "transform")
		.format_traceback(py)?;

		let output = transform
		.call1(
		py,
		PyTuple::new(
		py,
		&[task.into_py(py), args.into_py(py), inputs.into_py(py)],
		),
		)
		.format_traceback(py)?;

		output.extract(py).format_traceback(py)
		})?;

		Ok(serde_json::from_str(&results)?)
		}
		create_pymodule!("/src/bindings/transformers/transformers.py");

		pub fn get_model_from(task: &Value) -> Result<String> {
		Python::with_gil(\|py\| -> Result<String> {
Expand Down