From 5cb5a2b7cde301b2c1807a07cec8347e2a1c9c2b Mon Sep 17 00:00:00 2001
From: Wenxuan Zhang <wenxuangm@gmail.com>
Date: Fri, 21 Jan 2022 17:36:47 +0800
Subject: [PATCH] feat(sdk/py): implement push api

---
 sdk/python_v2/examples/push.py | 15 +++++++++++++++
 sdk/python_v2/src/convert.rs   | 25 +++++++++++++++++++++++--
 sdk/python_v2/src/lib.rs       | 27 +++++++++++++++++++++++++--
 3 files changed, 63 insertions(+), 4 deletions(-)
 create mode 100644 sdk/python_v2/examples/push.py

diff --git a/sdk/python_v2/examples/push.py b/sdk/python_v2/examples/push.py
new file mode 100644
index 000000000..95df04a5b
--- /dev/null
+++ b/sdk/python_v2/examples/push.py
@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+import asyncio
+from oomclient import Client
+
+
+async def main():
+    client = await Client.connect("http://localhost:50051")
+    kv_pairs = {
+        "last_5_click_posts": "1,3,5,7,9",
+        "number_of_user_starred_posts": 30,
+    }
+    await client.push("user-click", "929", kv_pairs)
+
+
+asyncio.run(main())
diff --git a/sdk/python_v2/src/convert.rs b/sdk/python_v2/src/convert.rs
index c451a0d8f..d1b549b1f 100644
--- a/sdk/python_v2/src/convert.rs
+++ b/sdk/python_v2/src/convert.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use oomclient::Value;
 use pyo3::{exceptions::PyException, prelude::*};
 
-pub fn value_to_py(value: Option<Value>, py: Python) -> PyObject {
+pub fn value_to_py(value: Option<&Value>, py: Python) -> PyObject {
     value
         .map(|value| match value {
             Value::Int64(v) => v.to_object(py),
@@ -18,7 +18,7 @@ pub fn value_to_py(value: Option<Value>, py: Python) -> PyObject {
 
 pub fn value_map_to_py(m: HashMap<String, Option<Value>>, py: Python) -> PyObject {
     m.into_iter()
-        .map(|(k, v)| (k, value_to_py(v, py)))
+        .map(|(k, v)| (k, value_to_py(v.as_ref(), py)))
         .collect::<HashMap<_, _>>()
         .into_py(py)
 }
@@ -26,3 +26,24 @@ pub fn value_map_to_py(m: HashMap<String, Option<Value>>, py: Python) -> PyObjec
 pub fn err_to_py(err: impl std::error::Error) -> PyErr {
     PyException::new_err(format!("{:?}", err))
 }
+
+pub fn py_to_value(obj: &PyAny) -> PyResult<Value> {
+    obj.extract::<ValueWrapper>().map(|wrapper| match wrapper {
+        ValueWrapper::Int64(v) => Value::Int64(v),
+        ValueWrapper::Double(v) => Value::Double(v),
+        ValueWrapper::String(v) => Value::String(v),
+        ValueWrapper::Bool(v) => Value::Bool(v),
+        ValueWrapper::UnixMilli(v) => Value::UnixMilli(v),
+        ValueWrapper::Bytes(v) => Value::Bytes(v),
+    })
+}
+
+#[derive(FromPyObject, Debug)]
+enum ValueWrapper {
+    Int64(i64),
+    Double(f64),
+    String(String),
+    Bool(bool),
+    UnixMilli(i64),
+    Bytes(Vec<u8>),
+}
diff --git a/sdk/python_v2/src/lib.rs b/sdk/python_v2/src/lib.rs
index 5c8f818a1..fe6251501 100644
--- a/sdk/python_v2/src/lib.rs
+++ b/sdk/python_v2/src/lib.rs
@@ -1,10 +1,13 @@
 mod convert;
 mod error;
 
-use convert::{err_to_py, value_map_to_py};
+use convert::{err_to_py, py_to_value, value_map_to_py};
 use error::Error;
 use oomclient::Client as OomClient;
-use pyo3::{prelude::*, types::PyType};
+use pyo3::{
+    prelude::*,
+    types::{PyDict, PyType},
+};
 use pyo3_asyncio::tokio::future_into_py;
 use std::collections::HashMap;
 
@@ -99,6 +102,26 @@ impl Client {
                 .map_err(err_to_py)
         })
     }
+
+    pub fn push<'p>(
+        &mut self,
+        py: Python<'p>,
+        group: String,
+        entity_key: String,
+        kv_pairs: &PyDict,
+    ) -> PyResult<&'p PyAny> {
+        let mut kvs = Vec::with_capacity(kv_pairs.len());
+        for (k, v) in kv_pairs {
+            let name = k.extract::<String>()?;
+            let value = py_to_value(v)?;
+            kvs.push((name, value));
+        }
+
+        let mut inner = OomClient::clone(&self.inner);
+        future_into_py(py, async move {
+            inner.push(entity_key, group, kvs).await.map_err(err_to_py)
+        })
+    }
 }
 
 /// OomClient python module implemented in Rust.