From 143dc2e911498b5c9d306b672a4cc7f9a1d438fc Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Mon, 21 Feb 2011 00:58:52 +0000
Subject: [PATCH] Added oid parameter to register_hstore()

The parameter is mostly useful with async connections that would need a
different protocol to be queried.

Issue reported by Jan "the Asynchronous".
---
 lib/extras.py         | 48 ++++++++++++++++++++++++++++---------------
 tests/types_extras.py | 20 ++++++++++++++++++
 2 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/lib/extras.py b/lib/extras.py
index 8e32c4f6..e1e94030 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -696,30 +696,44 @@ WHERE typname = 'hstore' and nspname = 'public';
 
         return oids
 
-def register_hstore(conn_or_curs, globally=False, unicode=False):
+def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
     """Register adapter and typecaster for `!dict`\-\ |hstore| conversions.
 
-    The function must receive a connection or cursor as the |hstore| oid is
-    different in each database. The typecaster will normally be registered
-    only on the connection or cursor passed as argument. If your application
-    uses a single database you can pass *globally*\=True to have the typecaster
-    registered on all the connections.
+    :param conn_or_curs: a connection or cursor: the typecaster will be
+        registered only on this object unless *globally* is set to `!True`
+    :param globally: register the adapter globally not only on *conn_or_curs*
+    :param unicode: if `!True`, keys and values returned from the database
+        will be `!unicode` instead of `!str`. The option is not available on
+        Python 3
+    :param oid: the OID of the |hstore| type if known. If not, it will be
+        queried on *conn_or_curs*
 
-    On Python 2, by default the returned dicts will have `!str` objects as keys and values:
-    use *unicode*\=True to return `!unicode` objects instead.  When adapting a
-    dictionary both `!str` and `!unicode` keys and values are handled (the
-    `unicode` values will be converted according to the current
-    `~connection.encoding`). The option is not available on Python 3.
+    The connection or cursor passed to the function will be used to query the
+    database and look for the OID of the |hstore| type (which may be different
+    across databases). If querying is not desirable (e.g. with
+    :ref:`asynchronous connections <async-support>`) you may specify it in the
+    *oid* parameter (it can be found using a query such as :sql:`SELECT
+    'hstore'::regtype::oid;`).
+
+    Note that, when passing a dictionary from Python to the database, both
+    strings and unicode keys and values are supported. Dictionaries returned
+    from the database have keys/values according to the *unicode* parameter.
 
     The |hstore| contrib module must be already installed in the database
     (executing the ``hstore.sql`` script in your ``contrib`` directory).
     Raise `~psycopg2.ProgrammingError` if the type is not found.
+
+    .. versionchanged:: 2.4
+        added the *oid* parameter.
     """
-    oids = HstoreAdapter.get_oids(conn_or_curs)
-    if oids is None:
-        raise psycopg2.ProgrammingError(
-            "hstore type not found in the database. "
-            "please install it from your 'contrib/hstore.sql' file")
+    if oid is None:
+        oid = HstoreAdapter.get_oids(conn_or_curs)
+        if oid is None:
+            raise psycopg2.ProgrammingError(
+                "hstore type not found in the database. "
+                "please install it from your 'contrib/hstore.sql' file")
+        else:
+            oid = oid[0]  # for the moment we don't have a HSTOREARRAY
 
     # create and register the typecaster
     if sys.version_info[0] < 3 and unicode:
@@ -727,7 +741,7 @@ def register_hstore(conn_or_curs, globally=False, unicode=False):
     else:
         cast = HstoreAdapter.parse
 
-    HSTORE = _ext.new_type((oids[0],), "HSTORE", cast)
+    HSTORE = _ext.new_type((oid,), "HSTORE", cast)
     _ext.register_type(HSTORE, not globally and conn_or_curs or None)
     _ext.register_adapter(dict, HstoreAdapter)
 
diff --git a/tests/types_extras.py b/tests/types_extras.py
index 88949298..85c3cfab 100755
--- a/tests/types_extras.py
+++ b/tests/types_extras.py
@@ -337,6 +337,26 @@ class HstoreTestCase(unittest.TestCase):
         ok({u''.join(ab): u''.join(ab)})
         ok(dict(zip(ab, ab)))
 
+    @skip_if_no_hstore
+    def test_oid(self):
+        cur = self.conn.cursor()
+        cur.execute("select 'hstore'::regtype::oid")
+        oid = cur.fetchone()[0]
+
+        # Note: None as conn_or_cursor is just for testing: not public
+        # interface and it may break in future.
+        from psycopg2.extras import register_hstore
+        register_hstore(None, globally=True, oid=oid)
+        try:
+            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'})
+
+        finally:
+            psycopg2.extensions.string_types.pop(oid)
+
 
 def skip_if_no_composite(f):
     def skip_if_no_composite_(self):