From 8b327807aefcb2df56902c94f249c4fe831fdfe1 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Audrius=20Ka=C5=BEukauskas?= <audrius@neutrino.lt>
Date: Thu, 15 Nov 2012 16:03:45 +0200
Subject: [PATCH] Register HStore adapter and typecaster in psycopg2 dialect

---
 .../dialects/postgresql/psycopg2.py           | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg2.py b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
index 700f76793a..05286ce203 100644
--- a/lib/sqlalchemy/dialects/postgresql/psycopg2.py
+++ b/lib/sqlalchemy/dialects/postgresql/psycopg2.py
@@ -147,6 +147,7 @@ from .base import PGDialect, PGCompiler, \
                                 PGIdentifierPreparer, PGExecutionContext, \
                                 ENUM, ARRAY, _DECIMAL_TYPES, _FLOAT_TYPES,\
                                 _INT_TYPES
+from .hstore import HSTORE
 
 
 logger = logging.getLogger('sqlalchemy.dialects.postgresql')
@@ -195,6 +196,13 @@ class _PGArray(ARRAY):
             self.item_type.convert_unicode = "force"
         # end Py2K
 
+class _PGHStore(HSTORE):
+    def bind_processor(self, dialect):
+        return None
+
+    def result_processor(self, dialect, coltype):
+        return None
+
 # When we're handed literal SQL, ensure it's a SELECT-query. Since
 # 8.3, combining cursors and "FOR UPDATE" has been fine.
 SERVER_SIDE_CURSOR_RE = re.compile(
@@ -282,6 +290,7 @@ class PGDialect_psycopg2(PGDialect):
             ENUM : _PGEnum, # needs force_unicode
             sqltypes.Enum : _PGEnum, # needs force_unicode
             ARRAY : _PGArray, # needs force_unicode
+            HSTORE : _PGHStore,
         }
     )
 
@@ -300,6 +309,16 @@ 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]
 
     @classmethod
     def dbapi(cls):
@@ -346,6 +365,13 @@ class PGDialect_psycopg2(PGDialect):
                 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 fns:
             def on_connect(conn):
                 for fn in fns:
-- 
2.47.3