From 6d441b6e03717ebc830dfabe63bca1eb659e3590 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 26 Sep 2010 22:11:06 +0100
Subject: [PATCH] Added hstore typecaster registration.

---
 lib/extras.py         | 43 +++++++++++++++++++++++++++++++++++++++++++
 tests/types_extras.py | 23 +++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/lib/extras.py b/lib/extras.py
index 7bfede5e..7dc64e58 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -567,5 +567,48 @@ class HstoreAdapter(object):
 
     parse = classmethod(parse)
 
+def register_hstore(conn_or_curs):
+    """Register adapter/typecaster for dict/hstore reading/writing.
+
+    The adapter must be registered on a connection or cursor as the hstore
+    oid is different in every database.
+
+    Raise `~psycopg2.ProgrammingError` if hstore is not installed in the
+    target database.
+    """
+    if hasattr(conn_or_curs, 'execute'):
+        conn = conn_or_curs.connection
+        curs = conn_or_curs
+    else:
+        conn = conn_or_curs
+        curs = conn_or_curs.cursor()
+
+    # Store the transaction status of the connection to revert it after use
+    conn_status = conn.status
+
+    # get the oid for the hstore
+    curs.execute("""\
+SELECT t.oid, typarray
+FROM pg_type t JOIN pg_namespace ns
+    ON typnamespace = ns.oid
+WHERE typname = 'hstore' and nspname = 'public';
+""")
+    oids = curs.fetchone()
+
+    # revert the status of the connection as before the command
+    if (conn_status != _ext.STATUS_IN_TRANSACTION
+    and conn.isolation_level != _ext.ISOLATION_LEVEL_AUTOCOMMIT):
+        conn.rollback()
+
+    if oids is None:
+        raise psycopg2.ProgrammingError(
+            "hstore type not found in the database. "
+            "please install it from your 'contrib/hstore.sql' file")
+
+    # create and register the typecaster
+    HSTORE = _ext.new_type((oids[0],), "HSTORE", HstoreAdapter.parse)
+    _ext.register_type(HSTORE, conn_or_curs)
+    _ext.register_adapter(dict, HstoreAdapter)
+
 
 __all__ = filter(lambda k: not k.startswith('_'), locals().keys())
diff --git a/tests/types_extras.py b/tests/types_extras.py
index 480312ba..26ad3123 100644
--- a/tests/types_extras.py
+++ b/tests/types_extras.py
@@ -186,6 +186,29 @@ class HstoreTestCase(unittest.TestCase):
         ko('"a=>"1"')
         ko('"a"=>"1", "b"=>NUL')
 
+    def test_register_conn(self):
+        from psycopg2.extras import register_hstore
+
+        register_hstore(self.conn)
+        cur = self.conn.cursor()
+        cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore")
+        t = cur.fetchone()
+        self.assert_(t[0] is None)
+        self.assertEqual(t[1], {})
+        self.assertEqual(t[2], {'a': 'b'})
+
+    def test_register_curs(self):
+        from psycopg2.extras import register_hstore
+
+        cur = self.conn.cursor()
+        register_hstore(cur)
+        cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore")
+        t = cur.fetchone()
+        self.assert_(t[0] is None)
+        self.assertEqual(t[1], {})
+        self.assertEqual(t[2], {'a': 'b'})
+
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)