Merge "Update dialect for pg8000 version 1.16.0"

6 files changed, 157 insertions, 81 deletions

diff --git a/lib/sqlalchemy/dialects/mysql/provision.py b/lib/sqlalchemy/dialects/mysql/provision.py
index b86056da6..bbe752d78 100644
--- a/lib/sqlalchemy/dialects/mysql/provision.py
+++ b/lib/sqlalchemy/dialects/mysql/provision.py
@@ -9,7 +9,7 @@ from ...testing.provision import temp_table_keyword_args
 
 
 @generate_driver_url.for_db("mysql", "mariadb")
-def generate_driver_url(url, driver):
+def generate_driver_url(url, driver, query):
     backend = url.get_backend_name()
 
     if backend == "mysql":
@@ -18,7 +18,10 @@ def generate_driver_url(url, driver):
             backend = "mariadb"
 
     new_url = copy.copy(url)
+    new_url.query = dict(new_url.query)
     new_url.drivername = "%s+%s" % (backend, driver)
+    new_url.query.update(query)
+
     try:
         new_url.get_dialect()
     except exc.NoSuchModuleError:
diff --git a/lib/sqlalchemy/dialects/postgresql/pg8000.py b/lib/sqlalchemy/dialects/postgresql/pg8000.py
index 57c8f5a9a..e08332a57 100644
--- a/lib/sqlalchemy/dialects/postgresql/pg8000.py
+++ b/lib/sqlalchemy/dialects/postgresql/pg8000.py
@@ -9,13 +9,11 @@ r"""
     :name: pg8000
     :dbapi: pg8000
     :connectstring: postgresql+pg8000://user:password@host:port/dbname[?key=value&key=value...]
-    :url: https://pythonhosted.org/pg8000/
+    :url: https://pypi.org/project/pg8000/
 
-.. note::
-
-    The pg8000 dialect is **not tested as part of SQLAlchemy's continuous
-    integration** and may have unresolved issues.  The recommended PostgreSQL
-    dialect is psycopg2.
+.. versionchanged:: 1.4  The pg8000 dialect has been updated for version
+   1.16.5 and higher, and is again part of SQLAlchemy's continuous integration
+   with full feature support.
 
 .. _pg8000_unicode:
 
@@ -56,9 +54,6 @@ of the :ref:`psycopg2 <psycopg2_isolation_level>` dialect:
 * ``SERIALIZABLE``
 * ``AUTOCOMMIT``
 
-.. versionadded:: 0.9.5 support for AUTOCOMMIT isolation level when using
-   pg8000.
-
 .. seealso::
 
     :ref:`postgresql_isolation_level`
@@ -74,12 +69,16 @@ from uuid import UUID as _python_UUID
 from .base import _DECIMAL_TYPES
 from .base import _FLOAT_TYPES
 from .base import _INT_TYPES
+from .base import ENUM
+from .base import INTERVAL
 from .base import PGCompiler
 from .base import PGDialect
 from .base import PGExecutionContext
 from .base import PGIdentifierPreparer
 from .base import UUID
 from .json import JSON
+from .json import JSONB
+from .json import JSONPathType
 from ... import exc
 from ... import processors
 from ... import types as sqltypes
@@ -125,6 +124,40 @@ class _PGJSON(JSON):
         else:
             return super(_PGJSON, self).result_processor(dialect, coltype)
 
+    def get_dbapi_type(self, dbapi):
+        return dbapi.JSON
+
+
+class _PGJSONB(JSONB):
+    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)
+
+    def get_dbapi_type(self, dbapi):
+        return dbapi.JSONB
+
+
+class _PGJSONIndexType(sqltypes.JSON.JSONIndexType):
+    def get_dbapi_type(self, dbapi):
+        raise NotImplementedError("should not be here")
+
+
+class _PGJSONIntIndexType(sqltypes.JSON.JSONIntIndexType):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.INTEGER
+
+
+class _PGJSONStrIndexType(sqltypes.JSON.JSONStrIndexType):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.STRING
+
+
+class _PGJSONPathType(JSONPathType):
+    def get_dbapi_type(self, dbapi):
+        return 1009
+
 
 class _PGUUID(UUID):
     def bind_processor(self, dialect):
@@ -148,8 +181,67 @@ class _PGUUID(UUID):
             return process
 
 
+class _PGEnum(ENUM):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.UNKNOWN
+
+
+class _PGInterval(INTERVAL):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.INTERVAL
+
+    @classmethod
+    def adapt_emulated_to_native(cls, interval, **kw):
+        return _PGInterval(precision=interval.second_precision)
+
+
+class _PGTimeStamp(sqltypes.DateTime):
+    def get_dbapi_type(self, dbapi):
+        if self.timezone:
+            # TIMESTAMPTZOID
+            return 1184
+        else:
+            # TIMESTAMPOID
+            return 1114
+
+
+class _PGTime(sqltypes.Time):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.TIME
+
+
+class _PGInteger(sqltypes.Integer):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.INTEGER
+
+
+class _PGSmallInteger(sqltypes.SmallInteger):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.INTEGER
+
+
+class _PGNullType(sqltypes.NullType):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.NULLTYPE
+
+
+class _PGBigInteger(sqltypes.BigInteger):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.BIGINTEGER
+
+
+class _PGBoolean(sqltypes.Boolean):
+    def get_dbapi_type(self, dbapi):
+        return dbapi.BOOLEAN
+
+
 class PGExecutionContext_pg8000(PGExecutionContext):
-    pass
+    def pre_exec(self):
+        if not self.compiled:
+            return
+
+        if self.dialect._dbapi_version > (1, 16, 0):
+            self.set_input_sizes()
 
 
 class PGCompiler_pg8000(PGCompiler):
@@ -160,20 +252,11 @@ class PGCompiler_pg8000(PGCompiler):
             + self.process(binary.right, **kw)
         )
 
-    def post_process_text(self, text):
-        if "%%" in text:
-            util.warn(
-                "The SQLAlchemy postgresql dialect "
-                "now automatically escapes '%' in text() "
-                "expressions to '%%'."
-            )
-        return text.replace("%", "%%")
-
 
 class PGIdentifierPreparer_pg8000(PGIdentifierPreparer):
-    def _escape_identifier(self, value):
-        value = value.replace(self.escape_quote, self.escape_to_quote)
-        return value.replace("%", "%%")
+    def __init__(self, *args, **kwargs):
+        PGIdentifierPreparer.__init__(self, *args, **kwargs)
+        self._double_percents = False
 
 
 class PGDialect_pg8000(PGDialect):
@@ -195,9 +278,23 @@ class PGDialect_pg8000(PGDialect):
         {
             sqltypes.Numeric: _PGNumericNoBind,
             sqltypes.Float: _PGNumeric,
-            JSON: _PGJSON,
             sqltypes.JSON: _PGJSON,
+            sqltypes.Boolean: _PGBoolean,
+            sqltypes.NullType: _PGNullType,
+            JSONB: _PGJSONB,
+            sqltypes.JSON.JSONPathType: _PGJSONPathType,
+            sqltypes.JSON.JSONIndexType: _PGJSONIndexType,
+            sqltypes.JSON.JSONIntIndexType: _PGJSONIntIndexType,
+            sqltypes.JSON.JSONStrIndexType: _PGJSONStrIndexType,
             UUID: _PGUUID,
+            sqltypes.Interval: _PGInterval,
+            INTERVAL: _PGInterval,
+            sqltypes.DateTime: _PGTimeStamp,
+            sqltypes.Time: _PGTime,
+            sqltypes.Integer: _PGInteger,
+            sqltypes.SmallInteger: _PGSmallInteger,
+            sqltypes.BigInteger: _PGBigInteger,
+            sqltypes.Enum: _PGEnum,
         },
     )
 
@@ -313,6 +410,17 @@ class PGDialect_pg8000(PGDialect):
 
             fns.append(on_connect)
 
+        if self._dbapi_version > (1, 16, 0) and self._json_deserializer:
+
+            def on_connect(conn):
+                # json
+                conn.register_in_adapter(114, self._json_deserializer)
+
+                # jsonb
+                conn.register_in_adapter(3802, self._json_deserializer)
+
+            fns.append(on_connect)
+
         if len(fns) > 0:
 
             def on_connect(conn):
diff --git a/lib/sqlalchemy/dialects/postgresql/provision.py b/lib/sqlalchemy/dialects/postgresql/provision.py
index eb82a411e..6c6dc4be6 100644
--- a/lib/sqlalchemy/dialects/postgresql/provision.py
+++ b/lib/sqlalchemy/dialects/postgresql/provision.py
@@ -1,30 +1,13 @@
-import copy
 import time
 
 from ... import exc
 from ... import text
 from ...testing.provision import create_db
 from ...testing.provision import drop_db
-from ...testing.provision import generate_driver_url
 from ...testing.provision import log
 from ...testing.provision import temp_table_keyword_args
 
 
-@generate_driver_url.for_db("postgresql")
-def generate_driver_url(url, driver):
-    new_url = copy.copy(url)
-    new_url.drivername = "postgresql+%s" % driver
-    if new_url.get_driver_name() == "asyncpg":
-        new_url.query = dict(new_url.query)
-        new_url.query["async_fallback"] = "true"
-    try:
-        new_url.get_dialect()
-    except exc.NoSuchModuleError:
-        return None
-    else:
-        return new_url
-
-
 @create_db.for_db("postgresql")
 def _pg_create_db(cfg, eng, ident):
     template_db = cfg.options.postgresql_templatedb
diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py
index fe74be823..f9fabbeed 100644
--- a/lib/sqlalchemy/testing/assertions.py
+++ b/lib/sqlalchemy/testing/assertions.py
@@ -346,12 +346,12 @@ def _expect_raises(except_cls, msg=None, check_context=False):
     assert success, "Callable did not raise an exception"
 
 
-def expect_raises(except_cls):
-    return _expect_raises(except_cls, check_context=True)
+def expect_raises(except_cls, check_context=True):
+    return _expect_raises(except_cls, check_context=check_context)
 
 
-def expect_raises_message(except_cls, msg):
-    return _expect_raises(except_cls, msg=msg, check_context=True)
+def expect_raises_message(except_cls, msg, check_context=True):
+    return _expect_raises(except_cls, msg=msg, check_context=check_context)
 
 
 class AssertsCompiledSQL(object):
diff --git a/lib/sqlalchemy/testing/provision.py b/lib/sqlalchemy/testing/provision.py
index 21bacfca2..094d1ea94 100644
--- a/lib/sqlalchemy/testing/provision.py
+++ b/lib/sqlalchemy/testing/provision.py
@@ -7,6 +7,7 @@ from . import engines
 from .. import exc
 from ..engine import url as sa_url
 from ..util import compat
+from ..util import parse_qsl
 
 log = logging.getLogger(__name__)
 
@@ -85,7 +86,7 @@ def generate_db_urls(db_urls, extra_drivers):
         --dburi postgresql://db1  \
         --dburi postgresql://db2  \
         --dburi postgresql://db2  \
-        --dbdriver=psycopg2 --dbdriver=asyncpg
+        --dbdriver=psycopg2 --dbdriver=asyncpg?async_fallback=true
 
     Noting that the default postgresql driver is psycopg2.  the output
     would be::
@@ -139,21 +140,34 @@ def _generate_driver_urls(url, extra_drivers):
     main_driver = url.get_driver_name()
     extra_drivers.discard(main_driver)
 
-    url = generate_driver_url(url, main_driver)
+    url = generate_driver_url(url, main_driver, {})
     yield str(url)
 
     for drv in list(extra_drivers):
-        new_url = generate_driver_url(url, drv)
+
+        if "?" in drv:
+
+            driver_only, query_str = drv.split("?", 1)
+
+            query = parse_qsl(query_str)
+        else:
+            driver_only = drv
+            query = {}
+
+        new_url = generate_driver_url(url, driver_only, query)
         if new_url:
             extra_drivers.remove(drv)
+
             yield str(new_url)
 
 
 @register.init
-def generate_driver_url(url, driver):
+def generate_driver_url(url, driver, query):
     backend = url.get_backend_name()
     new_url = copy.copy(url)
+    new_url.query = dict(new_url.query)
     new_url.drivername = "%s+%s" % (backend, driver)
+    new_url.query.update(query)
     try:
         new_url.get_dialect()
     except exc.NoSuchModuleError:
diff --git a/lib/sqlalchemy/testing/suite/test_types.py b/lib/sqlalchemy/testing/suite/test_types.py
index 5e6ac1eab..6a390231b 100644
--- a/lib/sqlalchemy/testing/suite/test_types.py
+++ b/lib/sqlalchemy/testing/suite/test_types.py
@@ -34,7 +34,6 @@ from ... import testing
 from ... import Text
 from ... import Time
 from ... import TIMESTAMP
-from ... import type_coerce
 from ... import TypeDecorator
 from ... import Unicode
 from ... import UnicodeText
@@ -1161,37 +1160,6 @@ class JSONStringCastIndexTest(_LiteralRoundTripFixture, fixtures.TablesTest):
             and_(name == "r6", cast(col["b"], String) == '"some value"'), "r6"
         )
 
-    def test_crit_against_int_basic(self):
-        name = self.tables.data_table.c.name
-        col = self.tables.data_table.c["data"]
-
-        self._test_index_criteria(
-            and_(name == "r6", cast(col["a"], String) == "5"), "r6"
-        )
-
-    def _dont_test_crit_against_string_coerce_type(self):
-        name = self.tables.data_table.c.name
-        col = self.tables.data_table.c["data"]
-
-        self._test_index_criteria(
-            and_(
-                name == "r6",
-                cast(col["b"], String) == type_coerce("some value", JSON),
-            ),
-            "r6",
-            test_literal=False,
-        )
-
-    def _dont_test_crit_against_int_coerce_type(self):
-        name = self.tables.data_table.c.name
-        col = self.tables.data_table.c["data"]
-
-        self._test_index_criteria(
-            and_(name == "r6", cast(col["a"], String) == type_coerce(5, JSON)),
-            "r6",
-            test_literal=False,
-        )
-
 
 __all__ = (
     "UnicodeVarcharTest",