Merge pull request #50 from qaspen-python/feature/support_enum

Supported ENUM PostgreSQL Type

Author: chandr-andr

Cargo.lock

@@ -1,5 +1,6 @@
 import datetime
 import uuid
+from enum import Enum
 from ipaddress import IPv4Address
 from typing import Any, Dict, List, Union
 
@@ -282,6 +283,10 @@ async def test_deserialization_composite_into_python(
     """Test that it's possible to deserialize custom postgresql type."""
     await psql_pool.execute("DROP TABLE IF EXISTS for_test")
     await psql_pool.execute("DROP TYPE IF EXISTS all_types")
+    await psql_pool.execute("DROP TYPE IF EXISTS inner_type")
+    await psql_pool.execute("DROP TYPE IF EXISTS enum_type")
+    await psql_pool.execute("CREATE TYPE enum_type AS ENUM ('sad', 'ok', 'happy')")
+    await psql_pool.execute("CREATE TYPE inner_type AS (inner_value VARCHAR, some_enum enum_type)")
     create_type_query = """
     CREATE type all_types AS (
         bytea_ BYTEA,
@@ -316,7 +321,9 @@ async def test_deserialization_composite_into_python(
         uuid_arr UUID ARRAY,
         inet_arr INET ARRAY,
         jsonb_arr JSONB ARRAY,
-        json_arr JSON ARRAY
+        json_arr JSON ARRAY,
+        test_inner_value inner_type,
+        test_enum_type enum_type
     )
     """
     create_table_query = """
@@ -329,8 +336,14 @@ async def test_deserialization_composite_into_python(
     await psql_pool.execute(
         querystring=create_table_query,
     )
+
+    class TestEnum(Enum):
+        OK = "ok"
+        SAD = "sad"
+        HAPPY = "happy"
+
     await psql_pool.execute(
-        querystring="INSERT INTO for_test VALUES (ROW($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32))",  # noqa: E501
+        querystring="INSERT INTO for_test VALUES (ROW($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, ROW($33, $34), $35))",  # noqa: E501
         parameters=[
             b"Bytes",
             "Some String",
@@ -394,9 +407,16 @@ async def test_deserialization_composite_into_python(
                     },
                 ),
             ],
+            "inner type value",
+            "happy",
+            TestEnum.OK,
         ],
     )
 
+    class ValidateModelForInnerValueType(BaseModel):
+        inner_value: str
+        some_enum: TestEnum
+
     class ValidateModelForCustomType(BaseModel):
         bytea_: List[int]
         varchar_: str
@@ -432,6 +452,9 @@ class ValidateModelForCustomType(BaseModel):
         jsonb_arr: List[Dict[str, List[Union[str, int, List[str]]]]]
         json_arr: List[Dict[str, List[Union[str, int, List[str]]]]]
 
+        test_inner_value: ValidateModelForInnerValueType
+        test_enum_type: TestEnum
+
     class TopLevelModel(BaseModel):
         custom_type: ValidateModelForCustomType
 
@@ -446,6 +469,41 @@ class TopLevelModel(BaseModel):
     assert isinstance(model_result[0], TopLevelModel)
 
 
+async def test_enum_type(psql_pool: ConnectionPool) -> None:
+    """Test that we can decode ENUM type from PostgreSQL."""
+
+    class TestEnum(Enum):
+        OK = "ok"
+        SAD = "sad"
+        HAPPY = "happy"
+
+    class TestStrEnum(str, Enum):
+        OK = "ok"
+        SAD = "sad"
+        HAPPY = "happy"
+
+    await psql_pool.execute("DROP TABLE IF EXISTS for_test")
+    await psql_pool.execute("DROP TYPE IF EXISTS mood")
+    await psql_pool.execute(
+        "CREATE TYPE mood AS ENUM ('sad', 'ok', 'happy')",
+    )
+    await psql_pool.execute(
+        "CREATE TABLE for_test (test_mood mood, test_mood2 mood)",
+    )
+
+    await psql_pool.execute(
+        querystring="INSERT INTO for_test VALUES ($1, $2)",
+        parameters=[TestEnum.HAPPY, TestEnum.OK],
+    )
+
+    qs_result = await psql_pool.execute(
+        "SELECT * FROM for_test",
+    )
+    assert qs_result.result()[0]["test_mood"] == TestEnum.HAPPY.value
+    assert qs_result.result()[0]["test_mood"] != TestEnum.HAPPY
+    assert qs_result.result()[0]["test_mood2"] == TestStrEnum.OK
+
+
 async def test_custom_type_as_parameter(
     psql_pool: ConnectionPool,
 ) -> None:

python/tests/test_value_converter.py

@@ -1,6 +1,6 @@
 use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
 use macaddr::{MacAddr6, MacAddr8};
-use postgres_types::{Field, FromSql, Kind};
+use postgres_types::{Field, FromSql, Kind, ToSql};
 use serde_json::{json, Map, Value};
 use std::{fmt::Debug, net::IpAddr};
 use uuid::Uuid;
@@ -15,7 +15,7 @@ use pyo3::{
     Bound, Py, PyAny, Python, ToPyObject,
 };
 use tokio_postgres::{
-    types::{to_sql_checked, ToSql, Type},
+    types::{to_sql_checked, Type},
     Column, Row,
 };
 
@@ -447,6 +447,15 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult<
         return Ok(PythonDTO::PyIpAddress(id_address));
     }
 
+    // It's used for Enum.
+    // If StrEnum is used on Python side,
+    // we simply stop at the `is_instance_of::<PyString>``.
+    if let Ok(value_attr) = parameter.getattr("value") {
+        if let Ok(possible_string) = value_attr.extract::<String>() {
+            return Ok(PythonDTO::PyString(possible_string));
+        }
+    }
+
     Err(RustPSQLDriverError::PyToRustValueConversionError(format!(
         "Can not covert you type {parameter} into inner one",
     )))
@@ -691,15 +700,12 @@ fn postgres_bytes_to_py(
 pub fn composite_postgres_to_py(
     py: Python<'_>,
     fields: &Vec<Field>,
-    buf: &[u8],
+    buf: &mut &[u8],
+    custom_decoders: &Option<Py<PyDict>>,
 ) -> RustPSQLDriverPyResult<Py<PyAny>> {
-    let mut vec_buf: Vec<u8> = vec![];
-    vec_buf.extend_from_slice(buf);
-    let mut buf: &[u8] = vec_buf.as_slice();
-
     let result_py_dict: Bound<'_, PyDict> = PyDict::new_bound(py);
 
-    let num_fields = postgres_types::private::read_be_i32(&mut buf).map_err(|err| {
+    let num_fields = postgres_types::private::read_be_i32(buf).map_err(|err| {
         RustPSQLDriverError::RustToPyValueConversionError(format!(
             "Cannot read bytes data from PostgreSQL: {err}"
         ))
@@ -713,66 +719,111 @@ pub fn composite_postgres_to_py(
     }
 
     for field in fields {
-        let oid = postgres_types::private::read_be_i32(&mut buf).map_err(|err| {
+        let oid = postgres_types::private::read_be_i32(buf).map_err(|err| {
             RustPSQLDriverError::RustToPyValueConversionError(format!(
                 "Cannot read bytes data from PostgreSQL: {err}"
             ))
         })? as u32;
+
         if oid != field.type_().oid() {
             return Err(RustPSQLDriverError::RustToPyValueConversionError(
                 "unexpected OID".into(),
             ));
         }
 
-        result_py_dict.set_item(
-            field.name(),
-            postgres_bytes_to_py(py, field.type_(), &mut buf, false)?.to_object(py),
-        )?;
+        match field.type_().kind() {
+            Kind::Simple | Kind::Array(_) => {
+                result_py_dict.set_item(
+                    field.name(),
+                    postgres_bytes_to_py(py, field.type_(), buf, false)?.to_object(py),
+                )?;
+            }
+            Kind::Enum(_) => {
+                result_py_dict.set_item(
+                    field.name(),
+                    postgres_bytes_to_py(py, &Type::VARCHAR, buf, false)?.to_object(py),
+                )?;
+            }
+            _ => {
+                let (_, tail) = buf.split_at(4_usize);
+                *buf = tail;
+                result_py_dict.set_item(
+                    field.name(),
+                    raw_bytes_data_process(py, buf, field.name(), field.type_(), custom_decoders)?
+                        .to_object(py),
+                )?;
+            }
+        }
     }
 
     Ok(result_py_dict.to_object(py))
 }
 
-/// Convert type from postgres to python type.
+/// Process raw bytes from `PostgreSQL`.
 ///
 /// # Errors
 ///
 /// May return Err Result if cannot convert postgres
 /// type into rust one.
-pub fn postgres_to_py(
+pub fn raw_bytes_data_process(
     py: Python<'_>,
-    row: &Row,
-    column: &Column,
-    column_i: usize,
+    raw_bytes_data: &mut &[u8],
+    column_name: &str,
+    column_type: &Type,
     custom_decoders: &Option<Py<PyDict>>,
 ) -> RustPSQLDriverPyResult<Py<PyAny>> {
-    let raw_bytes_data = row.col_buffer(column_i);
-
     if let Some(custom_decoders) = custom_decoders {
         let py_encoder_func = custom_decoders
             .bind(py)
-            .get_item(column.name().to_lowercase());
+            .get_item(column_name.to_lowercase());
 
         if let Ok(Some(py_encoder_func)) = py_encoder_func {
-            return Ok(py_encoder_func.call((raw_bytes_data,), None)?.unbind());
+            return Ok(py_encoder_func
+                .call((raw_bytes_data.to_vec(),), None)?
+                .unbind());
         }
     }
 
-    let column_type = column.type_();
-    match raw_bytes_data {
-        Some(mut raw_bytes_data) => match column_type.kind() {
-            Kind::Simple | Kind::Array(_) => {
-                postgres_bytes_to_py(py, column.type_(), &mut raw_bytes_data, true)
-            }
-            Kind::Composite(fields) => composite_postgres_to_py(py, fields, raw_bytes_data),
-            _ => Err(RustPSQLDriverError::RustToPyValueConversionError(
-                column.type_().to_string(),
-            )),
-        },
-        None => Ok(py.None()),
+    match column_type.kind() {
+        Kind::Simple | Kind::Array(_) => {
+            postgres_bytes_to_py(py, column_type, raw_bytes_data, true)
+        }
+        Kind::Composite(fields) => {
+            composite_postgres_to_py(py, fields, raw_bytes_data, custom_decoders)
+        }
+        Kind::Enum(_) => postgres_bytes_to_py(py, &Type::VARCHAR, raw_bytes_data, true),
+        _ => Err(RustPSQLDriverError::RustToPyValueConversionError(
+            column_type.to_string(),
+        )),
     }
 }
 
+/// Convert type from postgres to python type.
+///
+/// # Errors
+///
+/// May return Err Result if cannot convert postgres
+/// type into rust one.
+pub fn postgres_to_py(
+    py: Python<'_>,
+    row: &Row,
+    column: &Column,
+    column_i: usize,
+    custom_decoders: &Option<Py<PyDict>>,
+) -> RustPSQLDriverPyResult<Py<PyAny>> {
+    let raw_bytes_data = row.col_buffer(column_i);
+    if let Some(mut raw_bytes_data) = raw_bytes_data {
+        return raw_bytes_data_process(
+            py,
+            &mut raw_bytes_data,
+            column.name(),
+            column.type_(),
+            custom_decoders,
+        );
+    }
+    Ok(py.None())
+}
+
 /// Convert python List of Dict type or just Dict into serde `Value`.
 ///
 /// # Errors