Merge pull request #101 from ddimmich/master

Postgres 9.4 Jsonb support

4 files changed, 175 insertions, 5 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/__init__.py b/lib/sqlalchemy/dialects/postgresql/__init__.py
index d1c768186..1ed43b635 100644
--- a/lib/sqlalchemy/dialects/postgresql/__init__.py
+++ b/lib/sqlalchemy/dialects/postgresql/__init__.py
@@ -15,7 +15,7 @@ from .base import \
     TSVECTOR
 from .constraints import ExcludeConstraint
 from .hstore import HSTORE, hstore
-from .json import JSON, JSONElement
+from .json import JSON, JSONElement, JSONB
 from .ranges import INT4RANGE, INT8RANGE, NUMRANGE, DATERANGE, TSRANGE, \
     TSTZRANGE
 
@@ -25,5 +25,5 @@ __all__ = (
     'DOUBLE_PRECISION', 'TIMESTAMP', 'TIME', 'DATE', 'BYTEA', 'BOOLEAN',
     'INTERVAL', 'ARRAY', 'ENUM', 'dialect', 'Any', 'All', 'array', 'HSTORE',
     'hstore', 'INT4RANGE', 'INT8RANGE', 'NUMRANGE', 'DATERANGE',
-    'TSRANGE', 'TSTZRANGE', 'json', 'JSON', 'JSONElement'
+    'TSRANGE', 'TSTZRANGE', 'json', 'JSON', 'JSONB', 'JSONElement'
 )
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 35f335252..6c4031b01 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1384,6 +1384,9 @@ class PGTypeCompiler(compiler.GenericTypeCompiler):
     def visit_JSON(self, type_):
         return "JSON"
 
+    def visit_JSONB(self, type_):
+        return "JSONB"
+
     def visit_INT4RANGE(self, type_):
         return "INT4RANGE"
 
diff --git a/lib/sqlalchemy/dialects/postgresql/json.py b/lib/sqlalchemy/dialects/postgresql/json.py
index 2e29185e8..d19dbe118 100644
--- a/lib/sqlalchemy/dialects/postgresql/json.py
+++ b/lib/sqlalchemy/dialects/postgresql/json.py
@@ -14,7 +14,7 @@ from ... import sql
 from ...sql import elements
 from ... import util
 
-__all__ = ('JSON', 'JSONElement')
+__all__ = ('JSON', 'JSONElement', 'JSONB')
 
 
 class JSONElement(elements.BinaryExpression):
@@ -197,3 +197,123 @@ class JSON(sqltypes.TypeEngine):
 
 
 ischema_names['json'] = JSON
+
+
+
+class JSONB(JSON):
+    """Represent the Postgresql JSONB type.
+
+    The :class:`.JSONB` type stores arbitrary JSONB format data, e.g.::
+
+        data_table = Table('data_table', metadata,
+            Column('id', Integer, primary_key=True),
+            Column('data', JSONB)
+        )
+
+        with engine.connect() as conn:
+            conn.execute(
+                data_table.insert(),
+                data = {"key1": "value1", "key2": "value2"}
+            )
+
+    :class:`.JSONB` provides several operations:
+
+    * Index operations::
+
+        data_table.c.data['some key']
+
+    * Index operations returning text (required for text comparison)::
+
+        data_table.c.data['some key'].astext == 'some value'
+
+    * Index operations with a built-in CAST call::
+
+        data_table.c.data['some key'].cast(Integer) == 5
+
+    * Path index operations::
+
+        data_table.c.data[('key_1', 'key_2', ..., 'key_n')]
+
+    * Path index operations returning text (required for text comparison)::
+
+        data_table.c.data[('key_1', 'key_2', ..., 'key_n')].astext == 'some value'
+
+    Index operations return an instance of :class:`.JSONElement`, which represents
+    an expression such as ``column -> index``.  This element then defines
+    methods such as :attr:`.JSONElement.astext` and :meth:`.JSONElement.cast`
+    for setting up type behavior.
+
+    The :class:`.JSON` type, when used with the SQLAlchemy ORM, does not detect
+    in-place mutations to the structure.  In order to detect these, the
+    :mod:`sqlalchemy.ext.mutable` extension must be used.  This extension will
+    allow "in-place" changes to the datastructure to produce events which
+    will be detected by the unit of work.  See the example at :class:`.HSTORE`
+    for a simple example involving a dictionary.
+
+    Custom serializers and deserializers are specified at the dialect level,
+    that is using :func:`.create_engine`.  The reason for this is that when
+    using psycopg2, the DBAPI only allows serializers at the per-cursor
+    or per-connection level.   E.g.::
+
+        engine = create_engine("postgresql://scott:tiger@localhost/test",
+                                json_serializer=my_serialize_fn,
+                                json_deserializer=my_deserialize_fn
+                        )
+
+    When using the psycopg2 dialect, the json_deserializer is registered
+    against the database using ``psycopg2.extras.register_default_json``.
+
+    .. versionadded:: 0.9.7
+
+    """
+
+    __visit_name__ = 'JSONB'
+    hashable = False
+
+    class comparator_factory(sqltypes.Concatenable.Comparator):
+        """Define comparison operations for :class:`.JSON`."""
+
+        def __getitem__(self, other):
+            """Get the value at a given key."""
+
+            return JSONElement(self.expr, other)
+
+        def _adapt_expression(self, op, other_comparator):
+            # How does one do equality?? jsonb also has "=" eg. '[1,2,3]'::jsonb = '[1,2,3]'::jsonb
+            if isinstance(op, custom_op):
+                if op.opstring in ['?', '?&', '?|', '@>', '<@']:
+                    return op, sqltypes.Boolean
+                if op.opstring == '->':
+                    return op, sqltypes.Text
+            return sqltypes.Concatenable.Comparator.\
+                _adapt_expression(self, op, other_comparator)
+
+        def has_key(self, other):
+            """Boolean expression.  Test for presence of a key.  Note that the
+            key may be a SQLA expression.
+            """
+            return self.expr.op('?')(other)
+
+        def has_all(self, other):
+            """Boolean expression.  Test for presence of all keys in jsonb
+            """
+            return self.expr.op('?&')(other)
+
+        def has_any(self, other):
+            """Boolean expression.  Test for presence of any key in jsonb
+            """
+            return self.expr.op('?|')(other)
+
+        def contains(self, other, **kwargs):
+            """Boolean expression.  Test if keys (or array) are a superset of/contained
+            the keys of the argument jsonb expression.
+            """
+            return self.expr.op('@>')(other)
+
+        def contained_by(self, other):
+            """Boolean expression.  Test if keys are a proper subset of the
+            keys of the argument jsonb expression.
+            """
+            return self.expr.op('<@')(other)
+
+ischema_names['jsonb'] = JSONB
\ No newline at end of file
diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py
index d70a0a52f..87250d467 100644
--- a/test/dialect/postgresql/test_types.py
+++ b/test/dialect/postgresql/test_types.py
@@ -17,7 +17,7 @@ from sqlalchemy import exc, schema, types
 from sqlalchemy.dialects.postgresql import base as postgresql
 from sqlalchemy.dialects.postgresql import HSTORE, hstore, array, \
             INT4RANGE, INT8RANGE, NUMRANGE, DATERANGE, TSRANGE, TSTZRANGE, \
-            JSON
+            JSON, JSONB
 import decimal
 from sqlalchemy import util
 from sqlalchemy.testing.util import round_decimal
@@ -1982,7 +1982,6 @@ class JSONRoundTripTest(fixtures.TablesTest):
                 },
         )
 
-
     def test_unicode_round_trip_python(self):
         engine = self._non_native_engine()
         self._test_unicode_round_trip(engine)
@@ -1991,3 +1990,51 @@ class JSONRoundTripTest(fixtures.TablesTest):
     def test_unicode_round_trip_native(self):
         engine = testing.db
         self._test_unicode_round_trip(engine)
+
+class JSONBTest(JSONTest):
+    def setup(self):
+        metadata = MetaData()
+        self.test_table = Table('test_table', metadata,
+            Column('id', Integer, primary_key=True),
+            Column('test_column', JSONB)
+        )
+        self.jsoncol = self.test_table.c.test_column
+
+    #Note - add fixture data for arrays []
+
+    def test_where_has_key(self):
+        self._test_where(
+            # hide from 2to3
+            getattr(self.jsoncol, 'has_key')('data'),
+            "test_table.test_column ? %(test_column_1)s"
+        )
+
+    def test_where_has_all(self):
+        self._test_where(
+            self.jsoncol.has_all({'name': 'r1', 'data': {"k1": "r1v1", "k2": "r1v2"}}),
+            "test_table.test_column ?& %(test_column_1)s"
+        )
+
+    def test_where_has_any(self):
+        self._test_where(
+            self.jsoncol.has_any(postgresql.array(['name', 'data'])),
+            "test_table.test_column ?| ARRAY[%(param_1)s, %(param_2)s]"
+        )
+
+    def test_where_contains(self):
+        self._test_where(
+            self.jsoncol.contains({"k1": "r1v1"}),
+            "test_table.test_column @> %(test_column_1)s"
+        )
+
+    def test_where_contained_by(self):
+        self._test_where(
+            self.jsoncol.contained_by({'foo': '1', 'bar': None}),
+            "test_table.test_column <@ %(test_column_1)s"
+        )
+
+
+class JSONBRoundTripTest(JSONRoundTripTest):
+    __only_on__ = ('postgresql >= 9.4',)
+
+