Merge branch 'release/1.3.0'

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "scyllapy"
-version = "1.2.1"
+version = "1.3.0"
 edition = "2021"
 
 [lib]

README.md

@@ -324,6 +324,55 @@ async def execute(scylla: Scylla) -> None:
     )
 ```
 
+## User defined types
+
+We also support user defined types. You can pass them as a parameter to query.
+Or parse it as a model in response.
+
+Here's binding example. Imagine we have defined a type in scylla like this:
+
+```cql
+CREATE TYPE IF NOT EXISTS test (
+    id int,
+    name text
+);
+```
+
+Now we need to define a model for it in python.
+
+```python
+from dataclasses import dataclass
+from scyllapy.extra_types import ScyllaPyUDT
+
+@dataclass
+class TestUDT(ScyllaPyUDT):
+    # Always define fields in the same order as in scylla.
+    # Otherwise you will get an error, or wrong data.
+    id: int
+    name: str
+
+async def execute(scylla: Scylla) -> None:
+    await scylla.execute(
+        "INSERT INTO table(id, udt_col) VALUES (?, ?)",
+        [1, TestUDT(id=1, name="test")],
+    )
+
+```
+
+We also support pydantic based models. Decalre them like this:
+
+```python
+from pydantic import BaseModel
+from scyllapy.extra_types import ScyllaPyUDT
+
+
+class TestUDT(BaseModel, ScyllaPyUDT):
+    # Always define fields in the same order as in scylla.
+    # Otherwise you will get an error, or wrong data.
+    id: int
+    name: str
+
+```
 
 # Query building
 

python/scyllapy/extra_types.py

@@ -1,3 +1,48 @@
+import dataclasses
+from typing import Any, List
+
 from ._internal.extra_types import BigInt, Counter, Double, SmallInt, TinyInt, Unset
 
-__all__ = ("BigInt", "Counter", "Double", "SmallInt", "TinyInt", "Unset")
+try:
+    import pydantic
+except ImportError:
+    pydantic = None
+
+
+class ScyllaPyUDT:
+    """
+    Class for declaring UDT models.
+
+    This class is a mixin for models like dataclasses and pydantic models,
+    or classes that have `__slots__` attribute.
+
+    It can be further extended to support other model types.
+    """
+
+    def __dump_udt__(self) -> List[Any]:
+        """
+        Method to dump UDT models to a dict.
+
+        This method returns a list of values in the order of the UDT fields.
+        Because in the protocol, UDT fields should be sent in the same order as
+        they were declared.
+        """
+        if dataclasses.is_dataclass(self):
+            values = []
+            for field in dataclasses.fields(self):
+                values.append(getattr(self, field.name))
+            return values
+        if pydantic is not None and isinstance(self, pydantic.BaseModel):
+            values = []
+            for param in self.__class__.__signature__.parameters:
+                values.append(getattr(self, param))
+            return values
+        if hasattr(self, "__slots__"):
+            values = []
+            for slot in self.__slots__:
+                values.append(getattr(self, slot))
+            return values
+        raise ValueError("Unsupported model type")
+
+
+__all__ = ("BigInt", "Counter", "Double", "SmallInt", "TinyInt", "Unset", "ScyllaPyUDT")

python/tests/test_extra_types.py

@@ -1,3 +1,4 @@
+from dataclasses import asdict, dataclass
 from typing import Any
 
 import pytest
@@ -69,3 +70,80 @@ async def test_unset(scylla: Scylla) -> None:
         f"INSERT INTO {table_name}(id, name) VALUES (?, ?)",
         [1, extra_types.Unset()],
     )
+
+
+@pytest.mark.anyio
+async def test_udts(scylla: Scylla) -> None:
+    @dataclass
+    class TestUDT(extra_types.ScyllaPyUDT):
+        id: int
+        name: str
+
+    table_name = random_string(4)
+
+    udt_val = TestUDT(id=1, name="test")
+    await scylla.execute(f"CREATE TYPE test_udt{table_name} (id int, name text)")
+    await scylla.execute(
+        f"CREATE TABLE {table_name} "
+        f"(id INT PRIMARY KEY, udt_col frozen<test_udt{table_name}>)",
+    )
+    await scylla.execute(
+        f"INSERT INTO {table_name} (id, udt_col) VALUES (?, ?)",
+        [1, udt_val],
+    )
+
+    res = await scylla.execute(f"SELECT * FROM {table_name}")
+    assert res.all() == [{"id": 1, "udt_col": asdict(udt_val)}]
+
+
+@pytest.mark.anyio
+async def test_nested_udts(scylla: Scylla) -> None:
+    @dataclass
+    class NestedUDT(extra_types.ScyllaPyUDT):
+        one: int
+        two: str
+
+    @dataclass
+    class TestUDT(extra_types.ScyllaPyUDT):
+        id: int
+        name: str
+        nested: NestedUDT
+
+    table_name = random_string(4)
+
+    udt_val = TestUDT(id=1, name="test", nested=NestedUDT(one=1, two="2"))
+    await scylla.execute(f"CREATE TYPE nested_udt{table_name} (one int, two text)")
+    await scylla.execute(
+        f"CREATE TYPE test_udt{table_name} "
+        f"(id int, name text, nested frozen<nested_udt{table_name}>)",
+    )
+    await scylla.execute(
+        f"CREATE TABLE {table_name} "
+        f"(id INT PRIMARY KEY, udt_col frozen<test_udt{table_name}>)",
+    )
+    await scylla.execute(
+        f"INSERT INTO {table_name} (id, udt_col) VALUES (?, ?)",
+        [1, udt_val],
+    )
+
+    res = await scylla.execute(f"SELECT * FROM {table_name}")
+    assert res.all() == [{"id": 1, "udt_col": asdict(udt_val)}]
+
+
+@pytest.mark.parametrize(
+    ["typ", "val"],
+    [
+        ("BIGINT", 1),
+        ("TINYINT", 1),
+        ("SMALLINT", 1),
+        ("INT", 1),
+        ("FLOAT", 1.0),
+        ("DOUBLE", 1.0),
+    ],
+)
+@pytest.mark.anyio
+async def test_autocast_positional(scylla: Scylla, typ: str, val: Any) -> None:
+    table_name = random_string(4)
+    await scylla.execute(f"CREATE TABLE {table_name}(id INT PRIMARY KEY, val {typ})")
+    prepared = await scylla.prepare(f"INSERT INTO {table_name}(id, val) VALUES (?, ?)")
+    await scylla.execute(prepared, [1, val])

python/tests/test_parsing.py

@@ -0,0 +1,19 @@
+import pytest
+from tests.utils import random_string
+
+from scyllapy import Scylla
+
+
+@pytest.mark.anyio
+async def test_udt_parsing(scylla: Scylla) -> None:
+    table_name = random_string(4)
+    await scylla.execute(f"CREATE TYPE test_udt{table_name} (id int, name text)")
+    await scylla.execute(
+        f"CREATE TABLE {table_name} "
+        f"(id int PRIMARY KEY, udt_col frozen<test_udt{table_name}>)",
+    )
+    await scylla.execute(
+        f"INSERT INTO {table_name} (id, udt_col) VALUES (1, {{id: 1, name: 'test'}})",
+    )
+    res = await scylla.execute(f"SELECT * FROM {table_name}")
+    assert res.all() == [{"id": 1, "udt_col": {"id": 1, "name": "test"}}]

python/tests/test_queries.py

@@ -28,3 +28,34 @@ class TestDTO:
     res = await scylla.execute(f"SELECT id FROM {table_name}")
 
     assert res.all(as_class=TestDTO) == [TestDTO(id=42)]
+
+
+@pytest.mark.anyio
+async def test_udt_as_dataclass(scylla: Scylla) -> None:
+    @dataclass
+    class UDTType:
+        id: int
+        name: str
+
+    @dataclass
+    class TestDTO:
+        id: int
+        udt_col: UDTType
+
+        def __post_init__(self) -> None:
+            if not isinstance(self.udt_col, UDTType):
+                self.udt_col = UDTType(**self.udt_col)
+
+    table_name = random_string(4)
+    await scylla.execute(f"CREATE TYPE test_udt{table_name} (id int, name text)")
+    await scylla.execute(
+        f"CREATE TABLE {table_name} "
+        f"(id int PRIMARY KEY, udt_col frozen<test_udt{table_name}>)",
+    )
+    await scylla.execute(
+        f"INSERT INTO {table_name} (id, udt_col) VALUES (1, {{id: 1, name: 'test'}})",
+    )
+    res = await scylla.execute(f"SELECT * FROM {table_name}")
+    assert res.all(as_class=TestDTO) == [
+        TestDTO(id=1, udt_col=UDTType(id=1, name="test")),
+    ]

src/batches.rs

@@ -121,7 +121,7 @@ impl ScyllaPyInlineBatch {
         self.inner.append_statement(query);
         if let Some(passed_params) = values {
             self.values
-                .push(parse_python_query_params(Some(passed_params), false)?);
+                .push(parse_python_query_params(Some(passed_params), false, None)?);
         } else {
             self.values.push(SerializedValues::new());
         }

src/exceptions/rust_err.rs

@@ -44,6 +44,8 @@ pub enum ScyllaPyError {
     RowsDowncastError(String),
     #[error("Cannot parse value of column {0} as {1}.")]
     ValueDowncastError(String, &'static str),
+    #[error("Cannot downcast UDT {0} of column {1}. Reason: {2}.")]
+    UDTDowncastError(String, String, String),
     #[error("Query didn't suppose to return anything.")]
     NoReturnsError,
     #[error("Query doesn't have columns.")]
@@ -73,6 +75,7 @@ impl From<ScyllaPyError> for pyo3::PyErr {
             | ScyllaPyError::IpParseError(_) => ScyllaPyBindingError::new_err((err_desc,)),
             ScyllaPyError::RowsDowncastError(_)
             | ScyllaPyError::ValueDowncastError(_, _)
+            | ScyllaPyError::UDTDowncastError(_, _, _)
             | ScyllaPyError::NoReturnsError
             | ScyllaPyError::NoColumns => ScyllaPyMappingError::new_err((err_desc,)),
             ScyllaPyError::QueryBuilderError(_) => ScyllaPyQueryBuiderError::new_err((err_desc,)),

src/prepared_queries.rs

@@ -4,7 +4,7 @@ use scylla::prepared_statement::PreparedStatement;
 #[pyclass(name = "PreparedQuery")]
 #[derive(Clone, Debug)]
 pub struct ScyllaPyPreparedQuery {
-    inner: PreparedStatement,
+    pub inner: PreparedStatement,
 }
 
 impl From<PreparedStatement> for ScyllaPyPreparedQuery {

src/query_builder/delete.rs

@@ -110,7 +110,7 @@ impl Delete {
         slf.where_clauses_.push(clause);
         if let Some(vals) = values {
             for value in vals {
-                slf.values_.push(py_to_value(value)?);
+                slf.values_.push(py_to_value(value, None)?);
             }
         }
         Ok(slf)
@@ -148,7 +148,7 @@ impl Delete {
     ) -> ScyllaPyResult<PyRefMut<'a, Self>> {
         let parsed_values = if let Some(vals) = values {
             vals.iter()
-                .map(|item| py_to_value(item))
+                .map(|item| py_to_value(item, None))
                 .collect::<Result<Vec<_>, _>>()?
         } else {
             vec![]

src/query_builder/insert.rs

@@ -114,7 +114,7 @@ impl Insert {
         if value.is_none() {
             slf.values_.push(ScyllaPyCQLDTO::Unset);
         } else {
-            slf.values_.push(py_to_value(value)?);
+            slf.values_.push(py_to_value(value, None)?);
         }
         Ok(slf)
     }

src/query_builder/select.rs

@@ -153,7 +153,7 @@ impl Select {
         slf.where_clauses_.push(clause);
         if let Some(vals) = values {
             for value in vals {
-                slf.values_.push(py_to_value(value)?);
+                slf.values_.push(py_to_value(value, None)?);
             }
         }
         Ok(slf)

src/query_builder/update.rs

@@ -130,7 +130,7 @@ impl Update {
         value: &'a PyAny,
     ) -> ScyllaPyResult<PyRefMut<'a, Self>> {
         slf.assignments_.push(UpdateAssignment::Simple(name));
-        slf.values_.push(py_to_value(value)?);
+        slf.values_.push(py_to_value(value, None)?);
         Ok(slf)
     }
 
@@ -147,7 +147,7 @@ impl Update {
     ) -> ScyllaPyResult<PyRefMut<'a, Self>> {
         slf.assignments_
             .push(UpdateAssignment::Inc(name.clone(), name));
-        slf.values_.push(py_to_value(value)?);
+        slf.values_.push(py_to_value(value, None)?);
         Ok(slf)
     }
 
@@ -164,7 +164,7 @@ impl Update {
     ) -> ScyllaPyResult<PyRefMut<'a, Self>> {
         slf.assignments_
             .push(UpdateAssignment::Dec(name.clone(), name));
-        slf.values_.push(py_to_value(value)?);
+        slf.values_.push(py_to_value(value, None)?);
         Ok(slf)
     }
     /// Add where clause.
@@ -187,7 +187,7 @@ impl Update {
         slf.where_clauses_.push(clause);
         if let Some(vals) = values {
             for value in vals {
-                slf.where_values_.push(py_to_value(value)?);
+                slf.where_values_.push(py_to_value(value, None)?);
             }
         }
         Ok(slf)
@@ -248,7 +248,7 @@ impl Update {
     ) -> ScyllaPyResult<PyRefMut<'a, Self>> {
         let parsed_values = if let Some(vals) = values {
             vals.iter()
-                .map(|item| py_to_value(item))
+                .map(|item| py_to_value(item, None))
                 .collect::<Result<Vec<_>, _>>()?
         } else {
             vec![]

src/scylla_cls.rs

@@ -290,9 +290,13 @@ impl Scylla {
         params: Option<&'a PyAny>,
         paged: bool,
     ) -> ScyllaPyResult<&'a PyAny> {
+        let mut col_spec = None;
         // We need to prepare parameter we're going to use
         // in query.
-        let query_params = parse_python_query_params(params, true)?;
+        if let ExecuteInput::PreparedQuery(prepared) = &query {
+            col_spec = Some(prepared.inner.get_prepared_metadata().col_specs.as_ref());
+        }
+        let query_params = parse_python_query_params(params, true, col_spec)?;
         // We need this clone, to safely share the session between threads.
         let (query, prepared) = match query {
             ExecuteInput::Text(txt) => (Some(Query::new(txt)), None),
@@ -322,7 +326,11 @@ impl Scylla {
                 let mut batch_params = Vec::new();
                 if let Some(passed_params) = params {
                     for query_params in passed_params {
-                        batch_params.push(parse_python_query_params(Some(query_params), false)?);
+                        batch_params.push(parse_python_query_params(
+                            Some(query_params),
+                            false,
+                            None,
+                        )?);
                     }
                 }
                 (batch.into(), batch_params)