diff options
author | Amjith Ramanujam <amjith.r@gmail.com> | 2015-02-28 21:48:56 -0800 |
---|---|---|
committer | Amjith Ramanujam <amjith.r@gmail.com> | 2015-02-28 21:48:56 -0800 |
commit | 1e3a5a1338b6158ccca97d608898f79ed3c627ed (patch) | |
tree | 83115bc98891e11b62926260d409da61453eb7eb | |
parent | 0b8bfcf3daec11a0e277b4edafa19fd0feed0e29 (diff) | |
parent | 8e919f1b6c74859414fe472ffe4ede24c0a3f632 (diff) |
Merge pull request #165 from drocco007/master
Interpret incoming JSON as a string instead of via json.loads
-rw-r--r-- | pgcli/pgexecute.py | 49 | ||||
-rw-r--r-- | tests/test_pgexecute.py | 21 | ||||
-rw-r--r-- | tests/utils.py | 23 |
3 files changed, 85 insertions, 8 deletions
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index e669d512..151654e9 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -10,19 +10,46 @@ _logger = logging.getLogger(__name__) # Cast all database input to unicode automatically. # See http://initd.org/psycopg/docs/usage.html#unicode-handling for more info. -ext.register_type(psycopg2.extensions.UNICODE) -ext.register_type(psycopg2.extensions.UNICODEARRAY) +ext.register_type(ext.UNICODE) +ext.register_type(ext.UNICODEARRAY) ext.register_type(ext.new_type((705,), "UNKNOWN", ext.UNICODE)) ext.register_type(ext.new_type((51766,), "HSTORE", ext.UNICODE)) # Cast bytea fields to text. By default, this will render as hex strings with # Postgres 9+ and as escaped binary in earlier versions. -psycopg2.extensions.register_type( - psycopg2.extensions.new_type((17,), 'BYTEA_TEXT', psycopg2.STRING)) +ext.register_type(ext.new_type((17,), 'BYTEA_TEXT', psycopg2.STRING)) # When running a query, make pressing CTRL+C raise a KeyboardInterrupt # See http://initd.org/psycopg/articles/2014/07/20/cancelling-postgresql-statements-python/ -psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select) +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): @@ -105,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: |