uuid and boolean converter

Author: sinisaos

piccolo/engine/mysql.py

@@ -2,10 +2,13 @@
 
 import contextvars
 import json
+import uuid
 from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Mapping, Optional, Union
 
+from pymysql.constants import FIELD_TYPE
+from pymysql.converters import conversions
 from typing_extensions import Self
 
 from piccolo.engine.base import (
@@ -32,16 +35,45 @@
     from piccolo.table import Table
 
 
+# converters and formaters
 def backticks_format(querysting: str) -> str:
     return querysting.replace('"', "`")
 
 
-def converter(value: list) -> str:
+def convert_list(value: list) -> str:
     if isinstance(value, list):
         return json.dumps(value)
     return value
 
 
+def convert_bool(value: int) -> bool:
+    return bool(int(value)) if value is not None else None
+
+
+def convert_uuid(value: Any) -> Union[str, uuid.UUID]:
+    if isinstance(value, (bytes, bytearray)):
+        value = value.decode()
+    value = value.strip()
+    # check if string is uuid string
+    if len(value) == 36 and value.count("-") == 4:
+        try:
+            return uuid.UUID(value)
+        except ValueError:
+            return value
+    return value
+
+
+converters = conversions.copy()
+custom_decoders: dict[str, Any] = {
+    FIELD_TYPE.STRING: convert_uuid,
+    FIELD_TYPE.VAR_STRING: convert_uuid,
+    FIELD_TYPE.VARCHAR: convert_uuid,
+    FIELD_TYPE.CHAR: convert_uuid,
+    FIELD_TYPE.TINY: convert_bool,
+}
+converters.update(custom_decoders)
+
+
 @dataclass
 class AsyncBatch(BaseBatch):
     connection: Connection
@@ -282,6 +314,8 @@ def __init__(
         self.current_transaction = contextvars.ContextVar(
             f"mysql_current_transaction_{db_name}", default=None
         )
+        # converters
+        config["conv"] = converters
 
         super().__init__(
             engine_type="mysql",
@@ -392,7 +426,7 @@ async def _run_in_new_connection(
             args = []
         connection = await self.get_new_connection()
         # convert lists
-        params = tuple(converter(arg) for arg in args)
+        params = tuple(convert_list(arg) for arg in args)
         try:
             async with connection.cursor() as cursor:
                 await cursor.execute(query, params)

pyproject.toml

@@ -11,6 +11,7 @@ line_length = 79
 module = [
     "asyncpg.*",
     "aiomysql.*",
+    "pymysql.*",
     "colorama",
     "dateutil",
     "IPython",

tests/apps/fixtures/commands/test_load_dump_mysql.py

@@ -116,7 +116,7 @@ def _run_comparison(self, table_class_names: list[str]):
             {
                 "id": 1,
                 "bigint_col": 1,
-                "boolean_col": 1,
+                "boolean_col": True,
                 "bytea_col": b"hello",
                 "date_col": datetime.date(2021, 1, 1),
                 "foreignkey_col": 1,
@@ -131,7 +131,7 @@ def _run_comparison(self, table_class_names: list[str]):
                 "text_col": "hello",
                 "timestamp_col": datetime.datetime(2021, 1, 1, 0, 0),
                 "timestamptz_col": datetime.datetime(2021, 1, 1, 0, 0),
-                "uuid_col": "12783854-c012-4c15-8183-8eecb46f2c4e",
+                "uuid_col": uuid.UUID("12783854-c012-4c15-8183-8eecb46f2c4e"),
                 "varchar_col": "hello",
                 "unique_col": "hello",
                 "null_col": None,