Merge pull request #29 from Intreecom/feature/universal-execute

Intreecom · Sep 24, 2023 · b556a1e · b556a1e
2 parents 2efca8c + 11c7f7d
commit b556a1e
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 42 deletions.
diff --git a/src/query_builder/delete.rs b/src/query_builder/delete.rs
@@ -201,7 +201,7 @@ impl Delete {
         } else {
             self.values_.clone()
         };
-        scylla.native_execute(py, query, values, false)
+        scylla.native_execute(py, Some(query), None, values, false)
     }
 
     /// Add to batch

diff --git a/src/query_builder/insert.rs b/src/query_builder/insert.rs
@@ -165,7 +165,7 @@ impl Insert {
     pub fn execute<'a>(&'a self, py: Python<'a>, scylla: &'a Scylla) -> ScyllaPyResult<&'a PyAny> {
         let mut query = Query::new(self.build_query()?);
         self.request_params_.apply_to_query(&mut query);
-        scylla.native_execute(py, query, self.values_.clone(), false)
+        scylla.native_execute(py, Some(query), None, self.values_.clone(), false)
     }
 
     /// Add to batch

diff --git a/src/query_builder/select.rs b/src/query_builder/select.rs
@@ -248,7 +248,7 @@ impl Select {
     ) -> ScyllaPyResult<&'a PyAny> {
         let mut query = Query::new(self.build_query());
         self.request_params_.apply_to_query(&mut query);
-        scylla.native_execute(py, query, self.values_.clone(), paged)
+        scylla.native_execute(py, Some(query), None, self.values_.clone(), paged)
     }
 
     /// Add to batch

diff --git a/src/query_builder/update.rs b/src/query_builder/update.rs
@@ -285,7 +285,7 @@ impl Update {
         } else {
             values
         };
-        scylla.native_execute(py, query, values, false)
+        scylla.native_execute(py, Some(query), None, values, false)
     }
 
     /// Add to batch

diff --git a/src/scylla_cls.rs b/src/scylla_cls.rs
@@ -13,7 +13,7 @@ use openssl::{
     x509::X509,
 };
 use pyo3::{pyclass, pymethods, PyAny, Python};
-use scylla::{frame::value::ValueList, query::Query};
+use scylla::{frame::value::ValueList, prepared_statement::PreparedStatement, query::Query};
 
 #[pyclass(frozen, weakref)]
 #[derive(Clone)]
@@ -52,7 +52,8 @@ impl Scylla {
     pub fn native_execute<'a>(
         &'a self,
         py: Python<'a>,
-        query: impl Into<Query> + Send + 'static,
+        query: Option<impl Into<Query> + Send + 'static>,
+        prepared: Option<PreparedStatement>,
         values: impl ValueList + Send + 'static,
         paged: bool,
     ) -> ScyllaPyResult<&'a PyAny> {
@@ -64,13 +65,31 @@ impl Scylla {
             ))?;
             // let res = session.query(query, values).await?;
             if paged {
-                Ok(ScyllaPyQueryReturns::IterableQueryResult(
-                    ScyllaPyIterableQueryResult::new(session.query_iter(query, values).await?),
-                ))
+                match (query, prepared) {
+                    (Some(query), None) => Ok(ScyllaPyQueryReturns::IterableQueryResult(
+                        ScyllaPyIterableQueryResult::new(session.query_iter(query, values).await?),
+                    )),
+                    (None, Some(prepared)) => Ok(ScyllaPyQueryReturns::IterableQueryResult(
+                        ScyllaPyIterableQueryResult::new(
+                            session.execute_iter(prepared, values).await?,
+                        ),
+                    )),
+                    _ => Err(ScyllaPyError::SessionError(
+                        "You should pass either query or prepared query.".into(),
+                    )),
+                }
             } else {
-                Ok(ScyllaPyQueryReturns::QueryResult(ScyllaPyQueryResult::new(
-                    session.query(query, values).await?,
-                )))
+                match (query, prepared) {
+                    (Some(query), None) => Ok(ScyllaPyQueryReturns::QueryResult(
+                        ScyllaPyQueryResult::new(session.query(query, values).await?),
+                    )),
+                    (None, Some(prepared)) => Ok(ScyllaPyQueryReturns::QueryResult(
+                        ScyllaPyQueryResult::new(session.execute(&prepared, values).await?),
+                    )),
+                    _ => Err(ScyllaPyError::SessionError(
+                        "You should pass either query or prepared query.".into(),
+                    )),
+                }
             }
         })
         .map_err(Into::into)
@@ -275,36 +294,12 @@ impl Scylla {
         // in query.
         let query_params = parse_python_query_params(params, true)?;
         // We need this clone, to safely share the session between threads.
-        let session_arc = self.scylla_session.clone();
-        scyllapy_future(py, async move {
-            let session_guard = session_arc.read().await;
-            let session = session_guard.as_ref().ok_or(ScyllaPyError::SessionError(
-                "Session is not initialized.".into(),
-            ))?;
-            if paged {
-                let res = match query {
-                    ExecuteInput::Text(text) => session.query_iter(text, query_params).await?,
-                    ExecuteInput::Query(query) => session.query_iter(query, query_params).await?,
-                    ExecuteInput::PreparedQuery(prepared) => {
-                        session.execute_iter(prepared, query_params).await?
-                    }
-                };
-                Ok(ScyllaPyQueryReturns::IterableQueryResult(
-                    ScyllaPyIterableQueryResult::new(res),
-                ))
-            } else {
-                let res = match query {
-                    ExecuteInput::Text(text) => session.query(text, query_params).await?,
-                    ExecuteInput::Query(query) => session.query(query, query_params).await?,
-                    ExecuteInput::PreparedQuery(prepared) => {
-                        session.execute(&prepared.into(), query_params).await?
-                    }
-                };
-                Ok(ScyllaPyQueryReturns::QueryResult(ScyllaPyQueryResult::new(
-                    res,
-                )))
-            }
-        })
+        let (query, prepared) = match query {
+            ExecuteInput::Text(txt) => (Some(Query::new(txt)), None),
+            ExecuteInput::Query(query) => (Some(Query::from(query)), None),
+            ExecuteInput::PreparedQuery(prep) => (None, Some(PreparedStatement::from(prep))),
+        };
+        self.native_execute(py, query, prepared, query_params, paged)
     }
 
     /// Execute a batch statement.