mirror of
https://github.com/dbcli/pgcli
synced 2024-05-31 01:17:54 +00:00
Interpret incoming JSON as a string instead of via json.loads
Closes #163
This commit is contained in:
parent
072d605347
commit
8e919f1b6c
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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:
|
||||
|
|
Loading…
Reference in New Issue
Block a user