Add jsonb postgres support to pydal

pydal_jsonb.py

# -*- coding: utf-8 -*-

"""
    Adds some PostgreSQL jsonb support to pydal.
    Check https://www.postgresql.org/docs/9.6/static/functions-json.html for additional info.
    
    In order to use operators you must specify the 'jsonb' type on `Field` class.
    Included query operators:
        - `@>`: db.where(db.table.jsonb_field.jcontains({"foo": "bar"}))
        - `<@`: db.where(db.table.jsonb_field.jin({"foo": "bar"}))
        - `->`: db.where(db.table.jsonb_field.jget("foo") == "bar")
        - `->>`: db.where(db.table.jsonb_field.jgetv("foo") == "bar")
        - `#>`: db.where(db.table.jsonb_field.jpath("{foo,2}") == "bar")
        - `#>>`: db.where(db.table.jsonb_field.jpathv("{foo,2}") == "bar")
        - `?`: db.where(db.table.jsonb_field.jhas("foo"))
"""

from pydal.adapters import adapters
from pydal.adapters.postgres import PostgrePsycoBoolean
from pydal.dialects.postgre import PostgreDialectBooleanJSON
from pydal.dialects import dialects, register_expression, sqltype_for
from pydal.helpers.serializers import serializers
from pydal.objects import Expression, Query
from pydal.parsers.postgre import PostgreBooleanAutoJSONParser
from pydal.parsers import parsers, for_type as parse_type
from pydal.representers.postgre import PostgreArraysRepresenter
from pydal.representers import representers, for_type as repr_type


@adapters.register_for('postgres:psycopg2')
class JSONBPostgreAdapter(PostgrePsycoBoolean):
    def _get_json_dialect(self):
        return JSONBPostgreDialect

    def _get_json_parser(self):
        return JSONBPostgreParser


@dialects.register_for(JSONBPostgreAdapter)
class JSONBPostgreDialect(PostgreDialectBooleanJSON):
    @sqltype_for('jsonb')
    def type_jsonb(self):
        return 'jsonb'

    def _jcontains(self, field, data, query_env={}):
        return '(%s @> %s)' % (
            self.expand(field, query_env=query_env),
            self.expand(data, field.type, query_env=query_env))

    @register_expression('jcontains')
    def _jcontains_expr(self, expr, data):
        return Expression(expr.db, self._jcontains, expr, data, expr.type)

    def _jin(self, field, data, query_env={}):
        return '(%s <@ %s)' % (
            self.expand(field, query_env=query_env),
            self.expand(data, field.type, query_env=query_env))

    @register_expression('jin')
    def _jin_expr(self, expr, data):
        return Expression(expr.db, self._jin, expr, data, expr.type)

    def _jget(self, field, data, query_env={}):
        if not isinstance(data, int):
            data = self.expand(data, field_type='string', query_env=query_env)
        return '%s -> %s' % (self.expand(field, query_env=query_env), str(data))

    @register_expression('jget')
    def _jget_expr(self, expr, data):
        return Expression(expr.db, self._jget, expr, data, expr.type)

    def _jgetv(self, field, data, query_env={}):
        if not isinstance(data, int):
            data = self.expand(data, field_type='string', query_env=query_env)
        return '%s ->> %s' % (self.expand(field, query_env=query_env), str(data))

    @register_expression('jgetv')
    def _jgetv_expr(self, expr, data):
        return Expression(expr.db, self._jgetv, expr, data, 'string')

    def _jpath(self, field, data, query_env={}):
        return '%s #> %s' % (
            self.expand(field, query_env=query_env),
            self.expand(data, field_type='string', query_env=query_env))

    @register_expression('jpath')
    def _jpath_expr(self, expr, data):
        return Expression(expr.db, self._jpath, expr, data, expr.type)

    def _jpathv(self, field, data, query_env={}):
        return '%s #>> %s' % (
            self.expand(field, query_env=query_env),
            self.expand(data, field_type='string', query_env=query_env))

    @register_expression('jpathv')
    def _jpathv_expr(self, expr, data):
        return Expression(expr.db, self._jpathv, expr, data, 'string')

    def _jhas(self, field, data, query_env={}):
        return '%s ? %s' % (
            self.expand(field, query_env=query_env),
            self.expand(data, field_type='string', query_env=query_env))

    @register_expression('jhas')
    def _jhas_expr(self, expr, data):
        return Expression(expr.db, self._jhas, expr, data, 'boolean')


@parsers.register_for(JSONBPostgreAdapter)
class JSONBPostgreParser(PostgreBooleanAutoJSONParser):
    @parse_type('jsonb')
    def _jsonb(self, value):
        return value


@representers.register_for(JSONBPostgreAdapter)
class JSONBPostgreRepresenter(PostgreArraysRepresenter):
    @repr_type('jsonb')
    def _jsonb(self, value):
        return serializers.json(value)