- Add support for PostgreSQL with PyGreSQL

Change-Id: I040b75ff3b4110e7e8b26442a4eb226ba8c26715
Pull-request: https://github.com/zzzeek/sqlalchemy/pull/234

9 files changed, 285 insertions, 14 deletions

diff --git a/doc/build/changelog/changelog_11.rst b/doc/build/changelog/changelog_11.rst
index 53bd38a98..dc7d5105e 100644
--- a/doc/build/changelog/changelog_11.rst
+++ b/doc/build/changelog/changelog_11.rst
@@ -22,6 +22,12 @@
     :version: 1.1.0b1
 
     .. change::
+        :tags: feature, postgresql
+
+        Added a new dialect for the PyGreSQL Postgresql dialect.  Thanks
+        to Christoph Zwerschke and Kaolin Imago Fire for their efforts.
+
+    .. change::
         :tags: bug, orm
         :tickets: 3488
 
diff --git a/doc/build/changelog/migration_11.rst b/doc/build/changelog/migration_11.rst
index 6f0da3780..ac6cf1dc7 100644
--- a/doc/build/changelog/migration_11.rst
+++ b/doc/build/changelog/migration_11.rst
@@ -1935,6 +1935,15 @@ emits::
 
 :ticket:`2729`
 
+Support for PyGreSQL
+--------------------
+
+The `PyGreSQL <https://pypi.python.org/pypi/PyGreSQL>`_ DBAPI is now supported.
+
+.. seealso::
+
+    :ref:`dialect-postgresql-pygresql`
+
 The "postgres" module is removed
 ---------------------------------
 
diff --git a/doc/build/dialects/postgresql.rst b/doc/build/dialects/postgresql.rst
index 616924685..b4c90643d 100644
--- a/doc/build/dialects/postgresql.rst
+++ b/doc/build/dialects/postgresql.rst
@@ -182,28 +182,34 @@ For example::
       )
 
 psycopg2
---------------
+--------
 
 .. automodule:: sqlalchemy.dialects.postgresql.psycopg2
 
 pg8000
---------------
+------
 
 .. automodule:: sqlalchemy.dialects.postgresql.pg8000
 
 psycopg2cffi
---------------
+------------
 
 .. automodule:: sqlalchemy.dialects.postgresql.psycopg2cffi
 
 py-postgresql
---------------------
+-------------
 
 .. automodule:: sqlalchemy.dialects.postgresql.pypostgresql
 
+.. _dialect-postgresql-pygresql:
+
+pygresql
+--------
+
+.. automodule:: sqlalchemy.dialects.postgresql.pygresql
 
 zxjdbc
---------------
+------
 
 .. automodule:: sqlalchemy.dialects.postgresql.zxjdbc
 
diff --git a/lib/sqlalchemy/dialects/postgresql/__init__.py b/lib/sqlalchemy/dialects/postgresql/__init__.py
index 8aa4509be..ffd100f67 100644
--- a/lib/sqlalchemy/dialects/postgresql/__init__.py
+++ b/lib/sqlalchemy/dialects/postgresql/__init__.py
@@ -5,7 +5,8 @@
 # This module is part of SQLAlchemy and is released under
 # the MIT License: http://www.opensource.org/licenses/mit-license.php
 
-from . import base, psycopg2, pg8000, pypostgresql, zxjdbc, psycopg2cffi
+from . import base, psycopg2, pg8000, pypostgresql, pygresql, \
+    zxjdbc, psycopg2cffi
 
 base.dialect = psycopg2.dialect
 
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index eb3449e40..9d019b56e 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -2395,7 +2395,9 @@ class PGDialect(default.DefaultDialect):
                   i.relname
             """
 
-        t = sql.text(IDX_SQL, typemap={'attname': sqltypes.Unicode})
+        t = sql.text(IDX_SQL, typemap={
+            'relname': sqltypes.Unicode,
+            'attname': sqltypes.Unicode})
         c = connection.execute(t, table_oid=table_oid)
 
         indexes = defaultdict(lambda: defaultdict(dict))
diff --git a/lib/sqlalchemy/dialects/postgresql/pygresql.py b/lib/sqlalchemy/dialects/postgresql/pygresql.py
new file mode 100644
index 000000000..d30206613
--- /dev/null
+++ b/lib/sqlalchemy/dialects/postgresql/pygresql.py
@@ -0,0 +1,243 @@
+# postgresql/pygresql.py
+# Copyright (C) 2005-2016 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: http://www.opensource.org/licenses/mit-license.php
+
+"""
+.. dialect:: postgresql+pygresql
+    :name: pygresql
+    :dbapi: pgdb
+    :connectstring: postgresql+pygresql://user:password@host:port/dbname\
+[?key=value&key=value...]
+    :url: http://www.pygresql.org/
+"""
+
+import decimal
+import re
+
+from ... import exc, processors, util
+from ...types import Numeric, JSON as Json
+from ...sql.elements import Null
+from .base import PGDialect, PGCompiler, PGIdentifierPreparer, \
+    _DECIMAL_TYPES, _FLOAT_TYPES, _INT_TYPES, UUID
+from .hstore import HSTORE
+from .json import JSON, JSONB
+
+
+class _PGNumeric(Numeric):
+
+    def bind_processor(self, dialect):
+        return None
+
+    def result_processor(self, dialect, coltype):
+        if not isinstance(coltype, int):
+            coltype = coltype.oid
+        if self.asdecimal:
+            if coltype in _FLOAT_TYPES:
+                return processors.to_decimal_processor_factory(
+                    decimal.Decimal,
+                    self._effective_decimal_return_scale)
+            elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:
+                # PyGreSQL returns Decimal natively for 1700 (numeric)
+                return None
+            else:
+                raise exc.InvalidRequestError(
+                    "Unknown PG numeric type: %d" % coltype)
+        else:
+            if coltype in _FLOAT_TYPES:
+                # PyGreSQL returns float natively for 701 (float8)
+                return None
+            elif coltype in _DECIMAL_TYPES or coltype in _INT_TYPES:
+                return processors.to_float
+            else:
+                raise exc.InvalidRequestError(
+                    "Unknown PG numeric type: %d" % coltype)
+
+
+class _PGHStore(HSTORE):
+
+    def bind_processor(self, dialect):
+        if not dialect.has_native_hstore:
+            return super(_PGHStore, self).bind_processor(dialect)
+        hstore = dialect.dbapi.Hstore
+        def process(value):
+            if isinstance(value, dict):
+                return hstore(value)
+            return value
+        return process
+
+    def result_processor(self, dialect, coltype):
+        if not dialect.has_native_hstore:
+            return super(_PGHStore, self).result_processor(dialect, coltype)
+
+
+class _PGJSON(JSON):
+
+    def bind_processor(self, dialect):
+        if not dialect.has_native_json:
+            return super(_PGJSON, self).bind_processor(dialect)
+        json = dialect.dbapi.Json
+
+        def process(value):
+            if value is self.NULL:
+                value = None
+            elif isinstance(value, Null) or (
+                value is None and self.none_as_null):
+                return None
+            if value is None or isinstance(value, (dict, list)):
+                return json(value)
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        if not dialect.has_native_json:
+            return super(_PGJSON, self).result_processor(dialect, coltype)
+
+
+class _PGJSONB(JSONB):
+
+    def bind_processor(self, dialect):
+        if not dialect.has_native_json:
+            return super(_PGJSONB, self).bind_processor(dialect)
+        json = dialect.dbapi.Json
+
+        def process(value):
+            if value is self.NULL:
+                value = None
+            elif isinstance(value, Null) or (
+                value is None and self.none_as_null):
+                return None
+            if value is None or isinstance(value, (dict, list)):
+                return json(value)
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        if not dialect.has_native_json:
+            return super(_PGJSONB, self).result_processor(dialect, coltype)
+
+
+class _PGUUID(UUID):
+
+    def bind_processor(self, dialect):
+        if not dialect.has_native_uuid:
+            return super(_PGUUID, self).bind_processor(dialect)
+        uuid = dialect.dbapi.Uuid
+
+        def process(value):
+            if value is None:
+                return None
+            if isinstance(value, (str, bytes)):
+                if len(value) == 16:
+                    return uuid(bytes=value)
+                return uuid(value)
+            if isinstance(value, int):
+                return uuid(int=value)
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        if not dialect.has_native_uuid:
+            return super(_PGUUID, self).result_processor(dialect, coltype)
+        if not self.as_uuid:
+            def process(value):
+                if value is not None:
+                    return str(value)
+            return process
+
+
+class _PGCompiler(PGCompiler):
+
+    def visit_mod_binary(self, binary, operator, **kw):
+        return self.process(binary.left, **kw) + " %% " + \
+            self.process(binary.right, **kw)
+
+    def post_process_text(self, text):
+        return text.replace('%', '%%')
+
+
+class _PGIdentifierPreparer(PGIdentifierPreparer):
+
+    def _escape_identifier(self, value):
+        value = value.replace(self.escape_quote, self.escape_to_quote)
+        return value.replace('%', '%%')
+
+
+class PGDialect_pygresql(PGDialect):
+
+    driver = 'pygresql'
+
+    statement_compiler = _PGCompiler
+    preparer = _PGIdentifierPreparer
+
+    @classmethod
+    def dbapi(cls):
+        import pgdb
+        return pgdb
+
+    colspecs = util.update_copy(
+        PGDialect.colspecs,
+        {
+            Numeric: _PGNumeric,
+            HSTORE: _PGHStore,
+            Json: _PGJSON,
+            JSON: _PGJSON,
+            JSONB: _PGJSONB,
+            UUID: _PGUUID,
+        }
+    )
+
+    def __init__(self, **kwargs):
+        super(PGDialect_pygresql, self).__init__(**kwargs)
+        try:
+            version = self.dbapi.version
+            m = re.match(r'(\d+)\.(\d+)', version)
+            version = (int(m.group(1)), int(m.group(2)))
+        except (AttributeError, ValueError, TypeError):
+            version = (0, 0)
+        self.dbapi_version = version
+        if version < (5, 0):
+            has_native_hstore = has_native_json = has_native_uuid = False
+            if version != (0, 0):
+                util.warn("PyGreSQL is only fully supported by SQLAlchemy"
+                    " since version 5.0.")
+        else:
+            self.supports_unicode_statements = True
+            self.supports_unicode_binds = True
+            has_native_hstore = has_native_json = has_native_uuid = True
+        self.has_native_hstore = has_native_hstore
+        self.has_native_json = has_native_json
+        self.has_native_uuid = has_native_uuid
+
+    def create_connect_args(self, url):
+        opts = url.translate_connect_args(username='user')
+        if 'port' in opts:
+            opts['host'] = '%s:%s' % (
+                opts.get('host', '').rsplit(':', 1)[0], opts.pop('port'))
+        opts.update(url.query)
+        return [], opts
+
+    def is_disconnect(self, e, connection, cursor):
+        if isinstance(e, self.dbapi.Error):
+            if not connection:
+                return False
+            try:
+                connection = connection.connection
+            except AttributeError:
+                pass
+            else:
+                if not connection:
+                    return False
+            try:
+                return connection.closed
+            except AttributeError:  # PyGreSQL < 5.0
+                return connection._cnx is None
+        return False
+
+
+dialect = PGDialect_pygresql
diff --git a/test/dialect/postgresql/test_query.py b/test/dialect/postgresql/test_query.py
index 9f92a7830..c031e43de 100644
--- a/test/dialect/postgresql/test_query.py
+++ b/test/dialect/postgresql/test_query.py
@@ -761,6 +761,7 @@ class MatchTest(fixtures.TestBase, AssertsCompiledSQL):
 
     @testing.fails_on('postgresql+psycopg2', 'uses pyformat')
     @testing.fails_on('postgresql+pypostgresql', 'uses pyformat')
+    @testing.fails_on('postgresql+pygresql', 'uses pyformat')
     @testing.fails_on('postgresql+zxjdbc', 'uses qmark')
     @testing.fails_on('postgresql+psycopg2cffi', 'uses pyformat')
     def test_expression_positional(self):
diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py
index 8818a9941..6bcc4cf9a 100644
--- a/test/dialect/postgresql/test_types.py
+++ b/test/dialect/postgresql/test_types.py
@@ -544,11 +544,11 @@ class NumericInterpretationTest(fixtures.TestBase):
     __backend__ = True
 
     def test_numeric_codes(self):
-        from sqlalchemy.dialects.postgresql import psycopg2cffi, pg8000, \
-            psycopg2, base
+        from sqlalchemy.dialects.postgresql import pg8000, pygresql, \
+            psycopg2, psycopg2cffi, base
 
-        dialects = (pg8000.dialect(), psycopg2.dialect(),
-                    psycopg2cffi.dialect())
+        dialects = (pg8000.dialect(), pygresql.dialect(),
+                psycopg2.dialect(), psycopg2cffi.dialect())
         for dialect in dialects:
             typ = Numeric().dialect_impl(dialect)
             for code in base._INT_TYPES + base._FLOAT_TYPES + \
@@ -2757,7 +2757,10 @@ class JSONRoundTripTest(fixtures.TablesTest):
         result = engine.execute(
             select([data_table.c.data['k1'].astext])
         ).first()
-        assert isinstance(result[0], util.text_type)
+        if engine.dialect.returns_unicode_strings:
+            assert isinstance(result[0], util.text_type)
+        else:
+            assert isinstance(result[0], util.string_types)
 
     def test_query_returned_as_int(self):
         engine = testing.db
diff --git a/test/engine/test_execute.py b/test/engine/test_execute.py
index 76d60f207..b1c8673d1 100644
--- a/test/engine/test_execute.py
+++ b/test/engine/test_execute.py
@@ -177,8 +177,8 @@ class ExecuteTest(fixtures.TestBase):
         lambda: testing.against('mysql+mysqldb'), 'db-api flaky')
     @testing.fails_on_everything_except(
         'postgresql+psycopg2', 'postgresql+psycopg2cffi',
-        'postgresql+pypostgresql', 'mysql+mysqlconnector',
-        'mysql+pymysql', 'mysql+cymysql')
+        'postgresql+pypostgresql', 'postgresql+pygresql',
+        'mysql+mysqlconnector', 'mysql+pymysql', 'mysql+cymysql')
     def test_raw_python(self):
         def go(conn):
             conn.execute(