Merge remote-tracking branch 'origin/pr/132' into pr132

2 files changed, 83 insertions, 17 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/pg8000.py b/lib/sqlalchemy/dialects/postgresql/pg8000.py
index 4ccc90208..4bb376a96 100644
--- a/lib/sqlalchemy/dialects/postgresql/pg8000.py
+++ b/lib/sqlalchemy/dialects/postgresql/pg8000.py
@@ -13,17 +13,30 @@
 postgresql+pg8000://user:password@host:port/dbname[?key=value&key=value...]
     :url: https://pythonhosted.org/pg8000/
 
+
+.. _pg8000_unicode:
+
 Unicode
 -------
 
-When communicating with the server, pg8000 **always uses the server-side
-character set**.  SQLAlchemy has no ability to modify what character set
-pg8000 chooses to use, and additionally SQLAlchemy does no unicode conversion
-of any kind with the pg8000 backend. The origin of the client encoding setting
-is ultimately the CLIENT_ENCODING setting in postgresql.conf.
+pg8000 will encode / decode string values between it and the server using the
+PostgreSQL ``client_encoding`` parameter; by default this is the value in
+the ``postgresql.conf`` file, which often defaults to ``SQL_ASCII``.
+Typically, this can be changed to ``utf-8``, as a more useful default::
+
+    #client_encoding = sql_ascii # actually, defaults to database
+                                 # encoding
+    client_encoding = utf8
+
+The ``client_encoding`` can be overriden for a session by executing the SQL:
+
+SET CLIENT_ENCODING TO 'utf8';
+
+SQLAlchemy will execute this SQL on all new connections based on the value
+passed to :func:`.create_engine` using the ``client_encoding`` parameter::
 
-It is not necessary, though is also harmless, to pass the "encoding" parameter
-to :func:`.create_engine` when using pg8000.
+    engine = create_engine(
+        "postgresql+pg8000://user:pass@host/dbname", client_encoding='utf8')
 
 
 .. _pg8000_isolation_level:
@@ -58,6 +71,8 @@ from ... import types as sqltypes
 from .base import (
     PGDialect, PGCompiler, PGIdentifierPreparer, PGExecutionContext,
     _DECIMAL_TYPES, _FLOAT_TYPES, _INT_TYPES)
+import re
+from sqlalchemy.dialects.postgresql.json import JSON
 
 
 class _PGNumeric(sqltypes.Numeric):
@@ -88,6 +103,15 @@ class _PGNumericNoBind(_PGNumeric):
         return None
 
 
+class _PGJSON(JSON):
+
+    def result_processor(self, dialect, coltype):
+        if dialect._dbapi_version > (1, 10, 1):
+            return None  # Has native JSON
+        else:
+            return super(_PGJSON, self).result_processor(dialect, coltype)
+
+
 class PGExecutionContext_pg8000(PGExecutionContext):
     pass
 
@@ -129,20 +153,29 @@ class PGDialect_pg8000(PGDialect):
         PGDialect.colspecs,
         {
             sqltypes.Numeric: _PGNumericNoBind,
-            sqltypes.Float: _PGNumeric
+            sqltypes.Float: _PGNumeric,
+            JSON: _PGJSON,
         }
     )
 
+    def __init__(self, client_encoding=None, **kwargs):
+        PGDialect.__init__(self, **kwargs)
+        self.client_encoding = client_encoding
+
     def initialize(self, connection):
-        if self.dbapi and hasattr(self.dbapi, '__version__'):
-            self._dbapi_version = tuple([
-                int(x) for x in
-                self.dbapi.__version__.split(".")])
-        else:
-            self._dbapi_version = (99, 99, 99)
         self.supports_sane_multi_rowcount = self._dbapi_version >= (1, 9, 14)
         super(PGDialect_pg8000, self).initialize(connection)
 
+    @util.memoized_property
+    def _dbapi_version(self):
+        if self.dbapi and hasattr(self.dbapi, '__version__'):
+            return tuple(
+                [
+                    int(x) for x in re.findall(
+                        r'(\d+)(?:[-\.]?|$)', self.dbapi.__version__)])
+        else:
+            return (99, 99, 99)
+
     @classmethod
     def dbapi(cls):
         return __import__('pg8000')
@@ -181,6 +214,16 @@ class PGDialect_pg8000(PGDialect):
                 (level, self.name, ", ".join(self._isolation_lookup))
             )
 
+    def set_client_encoding(self, connection, client_encoding):
+        # adjust for ConnectionFairy possibly being present
+        if hasattr(connection, 'connection'):
+            connection = connection.connection
+
+        cursor = connection.cursor()
+        cursor.execute("SET CLIENT_ENCODING TO '" + client_encoding + "'")
+        cursor.execute("COMMIT")
+        cursor.close()
+
     def do_begin_twophase(self, connection, xid):
         connection.connection.tpc_begin((0, xid, ''))
 
@@ -198,4 +241,24 @@ class PGDialect_pg8000(PGDialect):
     def do_recover_twophase(self, connection):
         return [row[1] for row in connection.connection.tpc_recover()]
 
+    def on_connect(self):
+        fns = []
+        if self.client_encoding is not None:
+            def on_connect(conn):
+                self.set_client_encoding(conn, self.client_encoding)
+            fns.append(on_connect)
+
+        if self.isolation_level is not None:
+            def on_connect(conn):
+                self.set_isolation_level(conn, self.isolation_level)
+            fns.append(on_connect)
+
+        if len(fns) > 0:
+            def on_connect(conn):
+                for fn in fns:
+                    fn(conn)
+            return on_connect
+        else:
+            return None
+
 dialect = PGDialect_pg8000
diff --git a/test/dialect/postgresql/test_dialect.py b/test/dialect/postgresql/test_dialect.py
index 9f86aaa7a..bdd292fff 100644
--- a/test/dialect/postgresql/test_dialect.py
+++ b/test/dialect/postgresql/test_dialect.py
@@ -99,11 +99,13 @@ class MiscTest(fixtures.TestBase, AssertsExecutionResults, AssertsCompiledSQL):
         assert 'will create implicit sequence' in msgs
         assert 'will create implicit index' in msgs
 
-    @testing.only_on('postgresql+psycopg2', 'psycopg2-specific feature')
+    @testing.only_on(
+        ['postgresql+psycopg2', 'postgresql+pg8000'],
+        'psycopg2/pg8000-specific feature')
     @engines.close_open_connections
     def test_client_encoding(self):
         c = testing.db.connect()
-        current_encoding = c.connection.connection.encoding
+        current_encoding = c.execute("show client_encoding").fetchone()[0]
         c.close()
 
         # attempt to use an encoding that's not
@@ -115,7 +117,8 @@ class MiscTest(fixtures.TestBase, AssertsExecutionResults, AssertsCompiledSQL):
 
         e = engines.testing_engine(options={'client_encoding': test_encoding})
         c = e.connect()
-        eq_(c.connection.connection.encoding, test_encoding)
+        new_encoding = c.execute("show client_encoding").fetchone()[0]
+        eq_(new_encoding, test_encoding)
 
     @testing.only_on(
         ['postgresql+psycopg2', 'postgresql+pg8000',