adding Bytea column

Author: dantownsend

piccolo/columns/__init__.py

@@ -1,5 +1,6 @@
 from .column_types import (  # noqa: F401
     Boolean,
+    Bytea,
     Date,
     Decimal,
     Float,

piccolo/columns/base.py

@@ -437,6 +437,8 @@ def get_sql_value(self, value: t.Any) -> t.Any:
             output = str(value).lower()
         elif isinstance(value, datetime.datetime):
             output = f"'{value.isoformat().replace('T', '')}'"
+        elif isinstance(value, bytes):
+            output = f"'{value.hex()}'"
         else:
             output = value
 

piccolo/columns/column_types.py

@@ -1218,3 +1218,64 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
                 return f"{select_string} {self.json_operator}"
             else:
                 return f"{select_string} {self.json_operator} AS {self.alias}"
+
+
+###############################################################################
+
+
+class Bytea(Column):
+    """
+    Used for storing bytes.
+
+    **Example**
+
+    .. code-block:: python
+
+        class Token(Table):
+            token = Bytea(default=b'token123')
+
+        # Create
+        >>> Token(token=b'my-token').save().run_sync()
+
+        # Query
+        >>> Token.select(Token.token).run_sync()
+        {'token': b'my-token'}
+
+    """
+
+    @property
+    def column_type(self):
+        engine_type = self._meta.table._meta.db.engine_type
+        if engine_type == "postgres":
+            return "BYTEA"
+        elif engine_type == "sqlite":
+            return "BLOB"
+        raise Exception("Unrecognized engine type")
+
+    def __init__(
+        self,
+        default: t.Union[
+            bytes,
+            bytearray,
+            t.Callable[[], bytes],
+            t.Callable[[], bytearray],
+            None,
+        ] = b"",
+        **kwargs,
+    ) -> None:
+        self._validate_default(default, (bytes, bytearray, None))
+
+        if isinstance(default, bytearray):
+            default = bytes(default)
+
+        self.default = default
+        kwargs.update({"default": default})
+        super().__init__(**kwargs)
+
+
+class Blob(Bytea):
+    """
+    An alias for Bytea.
+    """
+
+    pass

tests/columns/test_bytea.py

@@ -0,0 +1,62 @@
+from unittest import TestCase
+
+from piccolo.table import Table
+from piccolo.columns.column_types import Bytea
+
+
+class MyTable(Table):
+    token = Bytea()
+
+
+class MyTableDefault(Table):
+    """
+    Test the different default types.
+    """
+
+    token = Bytea()
+    token_bytes = Bytea(default=b"my-token")
+    token_bytearray = Bytea(default=bytearray(b"my-token"))
+    token_none = Bytea(default=None, null=True)
+
+
+class TestBytea(TestCase):
+    def setUp(self):
+        MyTable.create_table().run_sync()
+
+    def tearDown(self):
+        MyTable.alter().drop_table().run_sync()
+
+    def test_bytea(self):
+        """
+        Test storing a valid bytes value.
+        """
+        row = MyTable(token=b"my-token")
+        row.save().run_sync()
+        self.assertEqual(row.token, b"my-token")
+
+        self.assertEqual(
+            MyTable.select(MyTable.token).first().run_sync(),
+            {"token": b"my-token"},
+        )
+
+
+class TestByteaDefault(TestCase):
+    def setUp(self):
+        MyTableDefault.create_table().run_sync()
+
+    def tearDown(self):
+        MyTableDefault.alter().drop_table().run_sync()
+
+    def test_json_default(self):
+        row = MyTableDefault()
+        row.save().run_sync()
+
+        self.assertEqual(row.token, b"")
+        self.assertEqual(row.token_bytes, b"my-token")
+        self.assertEqual(row.token_bytearray, b"my-token")
+        self.assertEqual(row.token_none, None)
+
+    def test_invalid_default(self):
+        with self.assertRaises(ValueError):
+            for value in ("a", 1, ("x", "y", "z")):
+                Bytea(default=value)