Interpret incoming JSON as a string instead of via json.loads

Closes #163

3 files changed, 81 insertions, 3 deletions

diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 41e47eb3..151654e9 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -23,6 +23,34 @@ ext.register_type(ext.new_type((17,), 'BYTEA_TEXT', psycopg2.STRING))
 # See http://initd.org/psycopg/articles/2014/07/20/cancelling-postgresql-statements-python/
 ext.set_wait_callback(psycopg2.extras.wait_select)
 
+
+def register_json_typecasters(conn, loads_fn):
+    """Set the function for converting JSON data for a connection.
+
+    Use the supplied function to decode JSON data returned from the database
+    via the given connection. The function should accept a single argument of
+    the data as a string encoded in the database's character encoding.
+    psycopg2's default handler for JSON data is json.loads.
+    http://initd.org/psycopg/docs/extras.html#json-adaptation
+
+    This function attempts to register the typecaster for both JSON and JSONB
+    types.
+
+    Returns a set that is a subset of {'json', 'jsonb'} indicating which types
+    (if any) were successfully registered.
+    """
+    available = set()
+
+    for name in ['json', 'jsonb']:
+        try:
+            psycopg2.extras.register_json(conn, loads=loads_fn, name=name)
+            available.add(name)
+        except psycopg2.ProgrammingError:
+            pass
+
+    return available
+
+
 class PGExecute(object):
 
     search_path_query = '''
@@ -104,6 +132,18 @@ class PGExecute(object):
             self.conn.close()
         self.conn = conn
         self.conn.autocommit = True
+        register_json_typecasters(self.conn, self._json_typecaster)
+
+    def _json_typecaster(self, json_data):
+        """Interpret incoming JSON data as a string.
+
+        The raw data is decoded using the connection's encoding, which defaults
+        to the database's encoding.
+
+        See http://initd.org/psycopg/docs/connection.html#connection.encoding
+        """
+
+        return json_data.decode(self.conn.encoding)
 
     def run(self, statement):
         """Execute the sql in the database and return the results. The results
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 94f98627..25e577c5 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -3,7 +3,7 @@
 import pytest
 import psycopg2
 from textwrap import dedent
-from utils import run, dbtest
+from utils import run, dbtest, requires_json, requires_jsonb
 
 @dbtest
 def test_conn(executor):
@@ -125,6 +125,25 @@ def test_bytea_field_support_in_output(executor):
 
     assert u'\\xdeadbeef' in run(executor, "select * from binarydata", join=True)
 
+
 @dbtest
 def test_unicode_support_in_unknown_type(executor):
     assert u'日本語' in run(executor, "SELECT '日本語' AS japanese;", join=True)
+
+
+@requires_json
+def test_json_renders_without_u_prefix(executor, expanded):
+    run(executor, "create table jsontest(d json)")
+    run(executor, """insert into jsontest (d) values ('{"name": "Éowyn"}')""")
+    result = run(executor, "SELECT d FROM jsontest LIMIT 1", join=True)
+
+    assert u'{"name": "Éowyn"}' in result
+
+
+@requires_jsonb
+def test_jsonb_renders_without_u_prefix(executor, expanded):
+    run(executor, "create table jsonbtest(d jsonb)")
+    run(executor, """insert into jsonbtest (d) values ('{"name": "Éowyn"}')""")
+    result = run(executor, "SELECT d FROM jsonbtest LIMIT 1", join=True)
+
+    assert u'{"name": "Éowyn"}' in result
diff --git a/tests/utils.py b/tests/utils.py
index 9b7376f1..b755f78e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,6 +1,8 @@
 import pytest
 import psycopg2
+import psycopg2.extras
 from pgcli.main import format_output
+from pgcli.pgexecute import register_json_typecasters
 
 # TODO: should this be somehow be divined from environment?
 POSTGRES_USER, POSTGRES_HOST = 'postgres', 'localhost'
@@ -11,17 +13,34 @@ def db_connection(dbname=None):
     conn.autocommit = True
     return conn
 
+
 try:
-    db_connection()
+    conn = db_connection()
     CAN_CONNECT_TO_DB = True
+    SERVER_VERSION = conn.server_version
+    json_types = register_json_typecasters(conn, lambda x: x)
+    JSON_AVAILABLE = 'json' in json_types
+    JSONB_AVAILABLE = 'jsonb' in json_types
 except:
-    CAN_CONNECT_TO_DB = False
+    CAN_CONNECT_TO_DB = JSON_AVAILABLE = JSONB_AVAILABLE = False
+    SERVER_VERSION = 0
+
 
 dbtest = pytest.mark.skipif(
     not CAN_CONNECT_TO_DB,
     reason="Need a postgres instance at localhost accessible by user 'postgres'")
 
 
+requires_json = pytest.mark.skipif(
+    not JSON_AVAILABLE,
+    reason='Postgres server unavailable or json type not defined')
+
+
+requires_jsonb = pytest.mark.skipif(
+    not JSONB_AVAILABLE,
+    reason='Postgres server unavailable or jsonb type not defined')
+
+
 def create_db(dbname):
     with db_connection().cursor() as cur:
         try: