Merge pull request #26 from qaspen-python/feature/add_macaddr_support

chandr-andr · web-flow · commit 6903eb5866a3 · 2024-03-28T20:02:23.000+01:00
Added support for MacAddr6 type
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -28,3 +28,4 @@ chrono-tz = "0.8.5"
 uuid = { version = "1.7.0", features = ["v4"] }
 serde_json = "1.0.113"
 futures-util = "0.3.30"
+macaddr = "1.0.1"
diff --git a/docs/examples/aiohttp/start_example.py b/docs/examples/aiohttp/start_example.py
@@ -3,6 +3,7 @@
 from typing import Any, cast
 
 from aiohttp import web
+
 from psqlpy import PSQLPool
 
 
diff --git a/docs/examples/fastapi/advanced_example.py b/docs/examples/fastapi/advanced_example.py
@@ -6,6 +6,7 @@
 import uvicorn
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
+
 from psqlpy import PSQLPool
 
 db_pool = PSQLPool(
diff --git a/docs/examples/fastapi/start_example.py b/docs/examples/fastapi/start_example.py
@@ -5,9 +5,10 @@
 import uvicorn
 from fastapi import Depends, FastAPI, Request
 from fastapi.responses import JSONResponse
-from psqlpy import Connection, PSQLPool
 from typing_extensions import Annotated
 
+from psqlpy import Connection, PSQLPool
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
@@ -26,7 +27,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
 
 async def db_connection(request: Request) -> Connection:
     """Retrieve new connection from connection pool and return it."""
-    return await (cast(PSQLPool, request.app.state.db_pool)).connection()
+    return await cast(PSQLPool, request.app.state.db_pool).connection()
 
 
 @app.get("/")
diff --git a/pyproject.toml b/pyproject.toml
@@ -53,7 +53,7 @@ profile = "black"
 multi_line_output = 3
 
 [tool.black]
-line-length = 79
+line-length = 89
 
 [tool.mypy]
 strict = true
@@ -113,7 +113,7 @@ ignore = [
 ]
 exclude = [".venv/"]
 mccabe = { max-complexity = 10 }
-line-length = 79
+line-length = 89
 
 [tool.ruff.per-file-ignores]
 "python/psqlpy/*" = ["PYI021"]
diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi
@@ -535,7 +535,7 @@ class Transaction:
             )
 
         ```
-        """  # noqa: E501
+        """
     async def savepoint(self: Self, savepoint_name: str) -> None:
         """Create new savepoint.
 
diff --git a/python/psqlpy/_internal/extra_types.pyi b/python/psqlpy/_internal/extra_types.pyi
@@ -62,3 +62,13 @@ class PyJSON:
         ### Parameters:
         - `value`: value for the JSON/JSONB field.
         """
+
+class PyMacAddr6:
+    """Represents MACADDR in PostgreSQL."""
+
+    def __init__(self, value: str) -> None:
+        """Construct new MacAddr.
+
+        ### Parameters:
+        - `value`: value for MACADDR field.
+        """
diff --git a/python/psqlpy/extra_types.py b/python/psqlpy/extra_types.py
@@ -1,9 +1,10 @@
-from ._internal.extra_types import BigInt, Integer, PyJSON, PyUUID, SmallInt
+from ._internal.extra_types import BigInt, Integer, PyJSON, PyMacAddr6, PyUUID, SmallInt
 
 __all__ = [
     "SmallInt",
     "Integer",
     "BigInt",
     "PyUUID",
     "PyJSON",
+    "PyMacAddr6",
 ]
diff --git a/python/tests/test_cursor.py b/python/tests/test_cursor.py
@@ -116,10 +116,7 @@ async def test_cursor_fetch_forward_all(
 
     rest_results = await test_cursor.fetch_forward_all()
 
-    assert (
-        len(rest_results.result())
-        == number_database_records - default_fetch_number
-    )
+    assert len(rest_results.result()) == number_database_records - default_fetch_number
 
 
 async def test_cursor_fetch_backward(
diff --git a/python/tests/test_transaction.py b/python/tests/test_transaction.py
@@ -242,11 +242,9 @@ async def test_transaction_fetch_row(
 ) -> None:
     connection = await psql_pool.connection()
     async with connection.transaction() as transaction:
-        database_single_query_result: typing.Final = (
-            await transaction.fetch_row(
-                f"SELECT * FROM  {table_name} LIMIT 1",
-                [],
-            )
+        database_single_query_result: typing.Final = await transaction.fetch_row(
+            f"SELECT * FROM  {table_name} LIMIT 1",
+            [],
         )
         result = database_single_query_result.result()
         assert isinstance(result, dict)
diff --git a/src/additional_types.rs b/src/additional_types.rs
@@ -0,0 +1,36 @@
+use macaddr::MacAddr6;
+use tokio_postgres::types::{FromSql, Type};
+
+#[derive(Debug)]
+pub struct RustMacAddr6 {
+    inner: MacAddr6,
+}
+
+impl RustMacAddr6 {
+    #[must_use]
+    pub fn new(inner: MacAddr6) -> Self {
+        RustMacAddr6 { inner }
+    }
+
+    #[must_use]
+    pub fn inner(&self) -> &MacAddr6 {
+        &self.inner
+    }
+}
+
+impl<'a> FromSql<'a> for RustMacAddr6 {
+    fn from_sql(
+        _ty: &Type,
+        raw: &'a [u8],
+    ) -> Result<RustMacAddr6, Box<dyn std::error::Error + Sync + Send>> {
+        if raw.len() == 6 {
+            let new_mac_address = MacAddr6::new(raw[0], raw[1], raw[2], raw[3], raw[4], raw[5]);
+            return Ok(RustMacAddr6::new(new_mac_address));
+        }
+        Err("Cannot convert PostgreSQL MACADDR into rust MacAddr6".into())
+    }
+
+    fn accepts(_ty: &Type) -> bool {
+        true
+    }
+}
diff --git a/src/exceptions/rust_errors.rs b/src/exceptions/rust_errors.rs
@@ -34,6 +34,8 @@ pub enum RustPSQLDriverError {
     DBEngineBuildError(#[from] deadpool_postgres::BuildError),
     #[error("Value convert has failed: {0}")]
     UUIDConvertError(#[from] uuid::Error),
+    #[error("Cannot convert provided string to MacAddr6")]
+    MacAddr6ConversionError(#[from] macaddr::ParseError),
 }
 
 impl From<RustPSQLDriverError> for pyo3::PyErr {
@@ -43,6 +45,7 @@ impl From<RustPSQLDriverError> for pyo3::PyErr {
             RustPSQLDriverError::PyError(err) => err,
             RustPSQLDriverError::DBEngineError(_)
             | RustPSQLDriverError::DBEnginePoolError(_)
+            | RustPSQLDriverError::MacAddr6ConversionError(_)
             | RustPSQLDriverError::DBEngineBuildError(_) => {
                 RustPSQLDriverPyBaseError::new_err((error_desc,))
             }
diff --git a/src/extra_types.rs b/src/extra_types.rs
@@ -1,5 +1,6 @@
 use std::str::FromStr;
 
+use macaddr::MacAddr6;
 use pyo3::{pyclass, pymethods, types::PyModule, PyAny, PyResult, Python};
 use serde_json::Value;
 use uuid::Uuid;
@@ -94,6 +95,30 @@ impl PyJSON {
     }
 }
 
+#[pyclass]
+#[derive(Clone)]
+pub struct PyMacAddr6 {
+    inner: MacAddr6,
+}
+
+impl PyMacAddr6 {
+    #[must_use]
+    pub fn inner(self) -> MacAddr6 {
+        self.inner
+    }
+}
+
+#[pymethods]
+impl PyMacAddr6 {
+    #[new]
+    #[allow(clippy::missing_errors_doc)]
+    pub fn new_macaddr6(value: &str) -> RustPSQLDriverPyResult<Self> {
+        Ok(Self {
+            inner: MacAddr6::from_str(value)?,
+        })
+    }
+}
+
 #[allow(clippy::module_name_repetitions)]
 #[allow(clippy::missing_errors_doc)]
 pub fn extra_types_module(_py: Python<'_>, pymod: &PyModule) -> PyResult<()> {
@@ -102,5 +127,6 @@ pub fn extra_types_module(_py: Python<'_>, pymod: &PyModule) -> PyResult<()> {
     pymod.add_class::<BigInt>()?;
     pymod.add_class::<PyUUID>()?;
     pymod.add_class::<PyJSON>()?;
+    pymod.add_class::<PyMacAddr6>()?;
     Ok(())
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,3 +1,4 @@
+pub mod additional_types;
 pub mod common;
 pub mod driver;
 pub mod exceptions;
diff --git a/src/value_converter.rs b/src/value_converter.rs
@@ -1,4 +1,5 @@
 use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
+use macaddr::MacAddr6;
 use serde_json::{json, Map, Value};
 use std::{fmt::Debug, net::IpAddr};
 use uuid::Uuid;
@@ -18,8 +19,9 @@ use tokio_postgres::{
 };
 
 use crate::{
+    additional_types::RustMacAddr6,
     exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
-    extra_types::{BigInt, Integer, PyJSON, PyUUID, SmallInt},
+    extra_types::{BigInt, Integer, PyJSON, PyMacAddr6, PyUUID, SmallInt},
 };
 
 pub type QueryParameter = (dyn ToSql + Sync);
@@ -51,6 +53,7 @@ pub enum PythonDTO {
     PyList(Vec<PythonDTO>),
     PyTuple(Vec<PythonDTO>),
     PyJson(Value),
+    PyMacAddr6(MacAddr6),
 }
 
 impl PythonDTO {
@@ -182,6 +185,9 @@ impl ToSql for PythonDTO {
             PythonDTO::PyIpAddress(pyidaddress) => {
                 <&IpAddr as ToSql>::to_sql(&pyidaddress, ty, out)?;
             }
+            PythonDTO::PyMacAddr6(pymacaddr) => {
+                <&[u8] as ToSql>::to_sql(&pymacaddr.as_bytes(), ty, out)?;
+            }
             PythonDTO::PyList(py_iterable) | PythonDTO::PyTuple(py_iterable) => {
                 let mut items = Vec::new();
                 for inner in py_iterable {
@@ -350,6 +356,12 @@ pub fn py_to_rust(parameter: &PyAny) -> RustPSQLDriverPyResult<PythonDTO> {
         ));
     }
 
+    if parameter.is_instance_of::<PyMacAddr6>() {
+        return Ok(PythonDTO::PyMacAddr6(
+            parameter.extract::<PyMacAddr6>()?.inner(),
+        ));
+    }
+
     if let Ok(id_address) = parameter.extract::<IpAddr>() {
         return Ok(PythonDTO::PyIpAddress(id_address));
     }
@@ -471,6 +483,7 @@ pub fn postgres_to_py(
         Type::INET_ARRAY => Ok(row
             .try_get::<_, Option<Vec<IpAddr>>>(column_i)?
             .to_object(py)),
+        // Convert JSON/JSONB into Serde Value, then into list or dict
         Type::JSONB | Type::JSON => {
             let db_json = row.try_get::<_, Option<Value>>(column_i)?;
 
@@ -479,6 +492,11 @@ pub fn postgres_to_py(
                 None => Ok(py.None().to_object(py)),
             }
         }
+        // Convert MACADDR into inner type for macaddr6, then into str
+        Type::MACADDR => {
+            let macaddr_ = row.try_get::<_, RustMacAddr6>(column_i)?;
+            Ok(macaddr_.inner().to_string().to_object(py))
+        }
         _ => Err(RustPSQLDriverError::RustToPyValueConversionError(
             column.type_().to_string(),
         )),

Original file line number	Diff line number	Diff line change
`@@ -535,7 +535,7 @@ class Transaction:`
`535`	`535`	`)`
`536`	`536`
`537`	`537`	```
`538`		`- """ # noqa: E501`
	`538`	`+ """`
`539`	`539`	`async def savepoint(self: Self, savepoint_name: str) -> None:`
`540`	`540`	`"""Create new savepoint.`
`541`	`541`