Skip to content

Commit

Permalink
Fixes #402, fix custom JSON/JSONB type support.
Browse files Browse the repository at this point in the history
  • Loading branch information
fantix committed Dec 8, 2018
1 parent be33ea8 commit 8947074
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 11 deletions.
13 changes: 13 additions & 0 deletions gino/dialects/asyncpg.py
Expand Up @@ -21,6 +21,9 @@

from . import base

JSON_COLTYPE = 114
JSONB_COLTYPE = 3802


class AsyncpgDBAPI(base.BaseDBAPI):
Error = asyncpg.PostgresError, asyncpg.InterfaceError
Expand Down Expand Up @@ -283,6 +286,15 @@ async def _on_metadata_drop_async(self, target, bind, checkfirst=False,
await self.drop_async(bind=bind, checkfirst=checkfirst)


class GinoNullType(sqltypes.NullType):
def result_processor(self, dialect, coltype):
if coltype == JSON_COLTYPE:
return JSON().result_processor(dialect, coltype)
if coltype == JSONB_COLTYPE:
return JSONB().result_processor(dialect, coltype)
return super().result_processor(dialect, coltype)


# noinspection PyAbstractClass
class AsyncpgDialect(PGDialect, base.AsyncDialectMixin):
driver = 'asyncpg'
Expand All @@ -299,6 +311,7 @@ class AsyncpgDialect(PGDialect, base.AsyncDialectMixin):
{
ENUM: AsyncEnum,
sqltypes.Enum: AsyncEnum,
sqltypes.NullType: GinoNullType,
}
)

Expand Down
10 changes: 0 additions & 10 deletions gino/dialects/base.py
Expand Up @@ -2,15 +2,12 @@
import weakref

from sqlalchemy import util
from sqlalchemy.dialects.postgresql import JSON, JSONB

# noinspection PyProtectedMember
from ..engine import _SAConnection, _SAEngine, _DBAPIConnection
from ..loader import Loader

DEFAULT = object()
JSON_COLTYPE = 114
JSONB_COLTYPE = 3802


class BaseDBAPI:
Expand Down Expand Up @@ -375,13 +372,6 @@ def _init_statement_prepared(cls, dialect, connection, dbapi_connection,
self.cursor = self.create_cursor()
return self

def get_result_processor(self, type_, colname, coltype):
if coltype == JSON_COLTYPE:
return JSON().result_processor(self.dialect, coltype)
if coltype == JSONB_COLTYPE:
return JSONB().result_processor(self.dialect, coltype)
return super().get_result_processor(type_, colname, coltype)


class AsyncDialectMixin:
cursor_cls = DBAPICursor
Expand Down
15 changes: 14 additions & 1 deletion tests/test_json.py
Expand Up @@ -176,13 +176,21 @@ class Test(db.Model):
age = db.IntegerProperty(default=18)


async def test_t291(bind):
async def test_t291_t402(bind):
from gino.dialects.asyncpg import JSON, JSONB

class CustomJSONB(db.TypeDecorator):
impl = JSONB

def process_result_value(self, *_):
return 123

class PropsTest(db.Model):
__tablename__ = 'props_test_291'
profile = db.Column(JSONB(), nullable=False, server_default='{}')
profile1 = db.Column(JSON(), nullable=False, server_default='{}')
profile2 = db.Column(CustomJSONB(),
nullable=False, server_default='{}')

bool = db.BooleanProperty()
bool1 = db.BooleanProperty(prop_name='profile1')
Expand All @@ -197,5 +205,10 @@ class PropsTest(db.Model):
assert isinstance(profile, dict)
profile1 = await bind.scalar('SELECT profile1 FROM props_test_291')
assert isinstance(profile1, dict)
profile2 = await bind.scalar('SELECT profile2 FROM props_test_291')
assert isinstance(profile2, dict)
custom_profile2 = await bind.scalar(PropsTest.select('profile2'))
assert isinstance(custom_profile2, int)
assert custom_profile2 == 123
finally:
await PropsTest.gino.drop()

0 comments on commit 8947074

Please sign in to comment.