Skip to content

Commit

Permalink
Fix JSON and enum type columns
Browse files Browse the repository at this point in the history
  • Loading branch information
pmdevita committed Feb 21, 2024
1 parent 1e40ad1 commit 6b0387d
Show file tree
Hide file tree
Showing 2 changed files with 79 additions and 5 deletions.
9 changes: 5 additions & 4 deletions databases/backends/common/records.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
import json
import enum
import typing
from datetime import date, datetime

from sqlalchemy.engine.interfaces import Dialect
from sqlalchemy.engine.row import Row as SQLRow
from sqlalchemy.sql.compiler import _CompileLabel
from sqlalchemy.sql.schema import Column
from sqlalchemy.sql.sqltypes import JSON
from sqlalchemy.types import TypeEngine

from databases.interfaces import Record as RecordInterface
Expand Down Expand Up @@ -63,10 +64,10 @@ def __getitem__(self, key: typing.Any) -> typing.Any:
processor = datatype._cached_result_processor(self._dialect, None)

if self._dialect.name not in DIALECT_EXCLUDE:
if isinstance(raw, dict):
raw = json.dumps(raw)
if isinstance(datatype, JSON):
return raw

if processor is not None and (not isinstance(raw, (datetime, date))):
if processor is not None and not isinstance(raw, (datetime, date, enum.Enum)):
return processor(raw)
return raw

Expand Down
75 changes: 74 additions & 1 deletion tests/test_databases.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import asyncio
import datetime
import decimal
import enum
import functools
import gc
import itertools
Expand Down Expand Up @@ -55,6 +56,30 @@ def process_result_value(self, value, dialect):
sqlalchemy.Column("published", sqlalchemy.DateTime),
)


class TshirtSize(enum.Enum):
SMALL = "SMALL"
MEDIUM = "MEDIUM"
LARGE = "LARGE"
XL = "XL"


class TshirtColor(enum.Enum):
BLUE = 0
GREEN = 1
YELLOW = 2
RED = 3


# Used to test Enum
tshirt_size = sqlalchemy.Table(
"tshirt_size",
metadata,
sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True),
sqlalchemy.Column("size", sqlalchemy.Enum(TshirtSize)),
sqlalchemy.Column("color", sqlalchemy.Enum(TshirtColor))
)

# Used to test JSON
session = sqlalchemy.Table(
"session",
Expand Down Expand Up @@ -957,7 +982,32 @@ async def test_decimal_field(database_url):

@pytest.mark.parametrize("database_url", DATABASE_URLS)
@async_adapter
async def test_json_field(database_url):
async def test_enum_field(database_url):
"""
Test JSON columns, to ensure correct cross-database support.
"""

async with Database(database_url) as database:
async with database.transaction(force_rollback=True):
# execute()
size = TshirtSize.SMALL
color = TshirtColor.GREEN
values = {"size": size, "color": color}
query = tshirt_size.insert()
await database.execute(query, values)

# fetch_all()
query = tshirt_size.select()
results = await database.fetch_all(query=query)

assert len(results) == 1
assert results[0]["size"] == size
assert results[0]["color"] == color


@pytest.mark.parametrize("database_url", DATABASE_URLS)
@async_adapter
async def test_json_dict_field(database_url):
"""
Test JSON columns, to ensure correct cross-database support.
"""
Expand All @@ -978,6 +1028,29 @@ async def test_json_field(database_url):
assert results[0]["data"] == {"text": "hello", "boolean": True, "int": 1}


@pytest.mark.parametrize("database_url", DATABASE_URLS)
@async_adapter
async def test_json_list_field(database_url):
"""
Test JSON columns, to ensure correct cross-database support.
"""

async with Database(database_url) as database:
async with database.transaction(force_rollback=True):
# execute()
data = ['lemon', 'raspberry', 'lime', 'pumice']
values = {"data": data}
query = session.insert()
await database.execute(query, values)

# fetch_all()
query = session.select()
results = await database.fetch_all(query=query)

assert len(results) == 1
assert results[0]["data"] == ['lemon', 'raspberry', 'lime', 'pumice']


@pytest.mark.parametrize("database_url", DATABASE_URLS)
@async_adapter
async def test_custom_field(database_url):
Expand Down

0 comments on commit 6b0387d

Please sign in to comment.