- hstore adjustments

3 files changed, 100 insertions, 53 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 625ece6a1..f1061c90b 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -968,6 +968,9 @@ class PGTypeCompiler(compiler.GenericTypeCompiler):
     def visit_BIGINT(self, type_):
         return "BIGINT"
 
+    def visit_HSTORE(self, type_):
+        return "HSTORE"
+
     def visit_datetime(self, type_):
         return self.visit_TIMESTAMP(type_)
 
diff --git a/lib/sqlalchemy/dialects/postgresql/hstore.py b/lib/sqlalchemy/dialects/postgresql/hstore.py
index 4797031fa..ee5dec168 100644
--- a/lib/sqlalchemy/dialects/postgresql/hstore.py
+++ b/lib/sqlalchemy/dialects/postgresql/hstore.py
@@ -11,7 +11,6 @@ from ... import types as sqltypes
 from ...sql import functions as sqlfunc
 from ...sql.operators import custom_op
 from ...exc import SQLAlchemyError
-from ...ext.mutable import Mutable
 
 __all__ = ('HStoreSyntaxError', 'HSTORE', 'hstore')
 
@@ -114,41 +113,56 @@ def _serialize_hstore(val):
                      for k, v in val.iteritems())
 
 
-class MutationDict(Mutable, dict):
-    def __setitem__(self, key, value):
-        """Detect dictionary set events and emit change events."""
-        dict.__setitem__(self, key, value)
-        self.changed()
+class HSTORE(sqltypes.Concatenable, sqltypes.TypeEngine):
+    """Represent the Postgresql HSTORE type.
 
-    def __delitem__(self, key, value):
-        """Detect dictionary del events and emit change events."""
-        dict.__delitem__(self, key, value)
-        self.changed()
+    The :class:`.HSTORE` type stores dictionaries containing strings, e.g.::
 
-    @classmethod
-    def coerce(cls, key, value):
-        """Convert plain dictionary to MutationDict."""
-        if not isinstance(value, MutationDict):
-            if isinstance(value, dict):
-                return MutationDict(value)
-            return Mutable.coerce(key, value)
-        else:
-            return value
+        data_table = Table('data_table', metadata,
+            Column('id', Integer, primary_key=True),
+            Column('data', HSTORE)
+        )
+
+        with engine.connect() as conn:
+            conn.execute(
+                data_table.insert(),
+                data = {"key1": "value1", "key2": "value2"}
+            )
+
+    :class:`.HSTORE` provides for a wide range of operations, including:
+
+    * :meth:`.HSTORE.comparatopr_factory.has_key`
+
+    * :meth:`.HSTORE.comparatopr_factory.has_all`
+
+    * :meth:`.HSTORE.comparatopr_factory.defined`
 
-    def __getstate__(self):
-        return dict(self)
+    For usage with the SQLAlchemy ORM, it may be desirable to combine
+    the usage of :class:`.HSTORE` with the :mod:`sqlalchemy.ext.mutable`
+    extension.  This extension will allow in-place changes to dictionary
+    values to be detected by the unit of work::
 
-    def __setstate__(self, state):
-        self.update(state)
+        from sqlalchemy.ext.mutable import Mutable
 
+        class MyClass(Base):
+            __tablename__ = 'data_table'
 
-class HSTORE(sqltypes.Concatenable, sqltypes.UserDefinedType):
-    """The column type for representing PostgreSQL's contrib/hstore type.  This
-    type is a miniature key-value store in a column.  It supports query
-    operators for all the usual operations on a map-like data structure.
+            id = Column(Integer, primary_key=True)
+            data = Column(Mutable.as_mutable(HSTORE))
+
+        my_object = session.query(MyClass).one()
+
+        # in-place mutation, requires Mutable extension
+        # in order for the ORM to detect
+        my_object.data['some_key'] = 'some value'
+
+        session.commit()
 
     """
-    class comparator_factory(sqltypes.UserDefinedType.Comparator):
+
+    __visit_name__ = 'HSTORE'
+
+    class comparator_factory(sqltypes.TypeEngine.Comparator):
         def has_key(self, other):
             """Boolean expression.  Test for presence of a key.  Note that the
             key may be a SQLA expression.
@@ -237,6 +251,15 @@ class HSTORE(sqltypes.Concatenable, sqltypes.UserDefinedType):
                     return op, sqltypes.Text
             return op, other_comparator.type
 
+    #@util.memoized_property
+    #@property
+    #def _expression_adaptations(self):
+    #    return {
+    #        operators.getitem: {
+    #           sqltypes.String: sqltypes.String
+    #        },
+    #    }
+
     def bind_processor(self, dialect):
         def process(value):
             if isinstance(value, dict):
@@ -245,9 +268,6 @@ class HSTORE(sqltypes.Concatenable, sqltypes.UserDefinedType):
                 return value
         return process
 
-    def get_col_spec(self):
-        return 'HSTORE'
-
     def result_processor(self, dialect, coltype):
         def process(value):
             if value is not None:
@@ -256,8 +276,6 @@ class HSTORE(sqltypes.Concatenable, sqltypes.UserDefinedType):
                 return value
         return process
 
-MutationDict.associate_with(HSTORE)
-
 
 class hstore(sqlfunc.GenericFunction):
     """Construct an hstore on the server side using the hstore function.
diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 05286ce20..73f712328 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -131,9 +131,16 @@ The psycopg2 dialect will log Postgresql NOTICE messages via the
     import logging
     logging.getLogger('sqlalchemy.dialects.postgresql').setLevel(logging.INFO)
 
+HSTORE type
+------------
 
-"""
+The psycopg2 dialect will make use of the
+``psycopg2.extensions.register_hstore()`` extension when using the HSTORE
+type.  This replaces SQLAlchemy's pure-Python HSTORE coercion which takes
+effect for other DBAPIs.
 
+"""
+from __future__ import absolute_import
 import re
 import logging
 
@@ -198,10 +205,16 @@ class _PGArray(ARRAY):
 
 class _PGHStore(HSTORE):
     def bind_processor(self, dialect):
-        return None
+        if dialect._has_native_hstore:
+            return None
+        else:
+            return super(_PGHStore, self).bind_processor(dialect)
 
     def result_processor(self, dialect, coltype):
-        return None
+        if dialect._has_native_hstore:
+            return None
+        else:
+            return super(_PGHStore, self).result_processor(dialect, coltype)
 
 # When we're handed literal SQL, ensure it's a SELECT-query. Since
 # 8.3, combining cursors and "FOR UPDATE" has been fine.
@@ -283,6 +296,8 @@ class PGDialect_psycopg2(PGDialect):
     preparer = PGIdentifierPreparer_psycopg2
     psycopg2_version = (0, 0)
 
+    _has_native_hstore = False
+
     colspecs = util.update_copy(
         PGDialect.colspecs,
         {
@@ -295,10 +310,13 @@ class PGDialect_psycopg2(PGDialect):
     )
 
     def __init__(self, server_side_cursors=False, use_native_unicode=True,
-                        client_encoding=None, **kwargs):
+                        client_encoding=None,
+                        use_native_hstore=True,
+                        **kwargs):
         PGDialect.__init__(self, **kwargs)
         self.server_side_cursors = server_side_cursors
         self.use_native_unicode = use_native_unicode
+        self.use_native_hstore = use_native_hstore
         self.supports_unicode_binds = use_native_unicode
         self.client_encoding = client_encoding
         if self.dbapi and hasattr(self.dbapi, '__version__'):
@@ -309,21 +327,18 @@ class PGDialect_psycopg2(PGDialect):
                                             int(x)
                                             for x in m.group(1, 2, 3)
                                             if x is not None)
-        self._hstore_oids = None
+
 
     def initialize(self, connection):
         super(PGDialect_psycopg2, self).initialize(connection)
-
-        if self.psycopg2_version >= (2, 4):
-            extras = __import__('psycopg2.extras').extras
-            oids = extras.HstoreAdapter.get_oids(connection.connection)
-            if oids is not None and oids[0]:
-                self._hstore_oids = oids[0], oids[1]
+        self._has_native_hstore = self.use_native_hstore and \
+                        self._hstore_oids(connection.connection) \
+                            is not None
 
     @classmethod
     def dbapi(cls):
-        psycopg = __import__('psycopg2')
-        return psycopg
+        import psycopg2
+        return psycopg2
 
     @util.memoized_property
     def _isolation_lookup(self):
@@ -348,6 +363,8 @@ class PGDialect_psycopg2(PGDialect):
         connection.set_isolation_level(level)
 
     def on_connect(self):
+        from psycopg2 import extras, extensions
+
         fns = []
         if self.client_encoding is not None:
             def on_connect(conn):
@@ -360,17 +377,17 @@ class PGDialect_psycopg2(PGDialect):
             fns.append(on_connect)
 
         if self.dbapi and self.use_native_unicode:
-            extensions = __import__('psycopg2.extensions').extensions
             def on_connect(conn):
                 extensions.register_type(extensions.UNICODE, conn)
             fns.append(on_connect)
 
-        extras = __import__('psycopg2.extras').extras
-        def on_connect(conn):
-            if self._hstore_oids is not None:
-                oid, array_oid = self._hstore_oids
-                extras.register_hstore(conn, oid=oid, array_oid=array_oid)
-        fns.append(on_connect)
+        if self.dbapi and self.use_native_hstore:
+            def on_connect(conn):
+                hstore_oids = self._hstore_oids(conn)
+                if hstore_oids is not None:
+                    oid, array_oid = hstore_oids
+                    extras.register_hstore(conn, oid=oid, array_oid=array_oid)
+            fns.append(on_connect)
 
         if fns:
             def on_connect(conn):
@@ -380,6 +397,15 @@ class PGDialect_psycopg2(PGDialect):
         else:
             return None
 
+    @util.memoized_instancemethod
+    def _hstore_oids(self, conn):
+        if self.psycopg2_version >= (2, 4):
+            from psycopg2 import extras
+            oids = extras.HstoreAdapter.get_oids(conn)
+            if oids is not None and oids[0]:
+                return oids[0:2]
+        return None
+
     def create_connect_args(self, url):
         opts = url.translate_connect_args(username='user')
         if 'port' in opts: