From 3dd211e90b3d7ec098f04054dcb64debc84bbc2c Mon Sep 17 00:00:00 2001
From: Evert Timberg <evert.timberg@gmail.com>
Date: Mon, 2 Aug 2021 18:19:13 -0400
Subject: [PATCH] feat: JSONB column

---
 src/ormar_postgres_extensions/__init__.py     |  5 +-
 .../fields/__init__.py                        |  1 +
 src/ormar_postgres_extensions/fields/jsonb.py | 19 ++++++++
 tests/fields/test_jsonb.py                    | 48 +++++++++++++++++++
 4 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 src/ormar_postgres_extensions/fields/jsonb.py
 create mode 100644 tests/fields/test_jsonb.py

diff --git a/src/ormar_postgres_extensions/__init__.py b/src/ormar_postgres_extensions/__init__.py
index 9eb8999..d8aba2f 100644
--- a/src/ormar_postgres_extensions/__init__.py
+++ b/src/ormar_postgres_extensions/__init__.py
@@ -1 +1,4 @@
-from .fields import PostgresUUID  # noqa: F401
+from .fields import (  # noqa: F401
+    PostgresJSONB,
+    PostgresUUID,
+)
diff --git a/src/ormar_postgres_extensions/fields/__init__.py b/src/ormar_postgres_extensions/fields/__init__.py
index a0cd7b1..cb31770 100644
--- a/src/ormar_postgres_extensions/fields/__init__.py
+++ b/src/ormar_postgres_extensions/fields/__init__.py
@@ -1 +1,2 @@
+from .jsonb import PostgresJSONB  # noqa: F401
 from .uuid import PostgresUUID  # noqa: F401
diff --git a/src/ormar_postgres_extensions/fields/jsonb.py b/src/ormar_postgres_extensions/fields/jsonb.py
new file mode 100644
index 0000000..be234f1
--- /dev/null
+++ b/src/ormar_postgres_extensions/fields/jsonb.py
@@ -0,0 +1,19 @@
+from typing import Any
+
+import ormar
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.types import TypeDecorator
+
+
+class PostgresJSONBTypeDecorator(TypeDecorator):
+    impl = postgresql.JSONB
+
+
+class PostgresJSONB(ormar.JSON):
+    """
+    Custom JSON field uses a native PG JSONB type
+    """
+
+    @classmethod
+    def get_column_type(cls, **kwargs: Any) -> PostgresJSONBTypeDecorator:
+        return PostgresJSONBTypeDecorator()
diff --git a/tests/fields/test_jsonb.py b/tests/fields/test_jsonb.py
new file mode 100644
index 0000000..54a37a5
--- /dev/null
+++ b/tests/fields/test_jsonb.py
@@ -0,0 +1,48 @@
+import json
+from typing import Optional
+
+import ormar
+import pytest
+
+from ormar_postgres_extensions.fields import PostgresJSONB
+from tests.database import (
+    database,
+    metadata,
+)
+
+
+class JSONBTestModel(ormar.Model):
+    class Meta:
+        database = database
+        metadata = metadata
+
+    id: int = ormar.Integer(primary_key=True)
+    data: dict = PostgresJSONB()
+
+
+class NullableJSONBTestModel(ormar.Model):
+    class Meta:
+        database = database
+        metadata = metadata
+
+    id: int = ormar.Integer(primary_key=True)
+    data: Optional[dict] = PostgresJSONB(nullable=True)
+
+
+@pytest.mark.asyncio
+async def test_create_model_with_jsonb(db):
+    created = await JSONBTestModel(data=json.dumps(dict(foo="bar"))).save()
+    assert created.data == {"foo": "bar"}
+
+    # Confirm the model got saved to the DB by querying it back
+    found = await JSONBTestModel.objects.get()
+    assert found.data == {"foo": "bar"}
+
+
+@pytest.mark.asyncio
+async def test_create_model_with_nullable_jsonb(db):
+    created = await NullableJSONBTestModel().save()
+    assert created.data is None
+
+    found = await NullableJSONBTestModel.objects.get()
+    assert found.data is None