X-Git-Url: https://git.saurik.com/cyql.git/blobdiff_plain/35e48155e3bf0f206d8483a9bd09fb61a2beeb21..refs/heads/master:/__init__.py

diff --git a/__init__.py b/__init__.py
index f4e6b9f..d40d9b1 100644
--- a/__init__.py
+++ b/__init__.py
@@ -17,14 +17,32 @@ import psycopg2.pool
 psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
 psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
 
+class ConnectionError():
+    pass
+
 class connect(object):
     def __init__(self, dsn):
+        options = dsn.copy()
+        if 'cache' in options:
+            del options['cache']
+
+        if 'cache' in dsn:
+            cached = True
+            cache = dsn['cache']
+        else:
+            cached = False
+            cache = {
+                'hstore': None,
+            }
+
         attempt = 0
         while True:
             try:
-                self.driver = psycopg2.connect(**dsn)
+                self.driver = psycopg2.connect(**options)
                 break
             except psycopg2.OperationalError, e:
+                if e.message.startswith('could not connect to server: '):
+                    raise ConnectionError()
                 if attempt == 2:
                     raise
                 attempt = attempt + 1
@@ -38,10 +56,21 @@ class connect(object):
         #    self.driver.close()
         #    raise
 
-        try:
-            psycopg2.extras.register_hstore(self.driver, globally=False, unicode=True)
-        except psycopg2.ProgrammingError, e:
-            pass
+        hstore = cache['hstore']
+        if hstore == None:
+            hstore = psycopg2.extras.HstoreAdapter.get_oids(self.driver)
+            if hstore != None:
+                hstore = hstore[0]
+                cache['hstore'] = hstore
+
+        if hstore != None:
+            try:
+                psycopg2.extras.register_hstore(self.driver, globally=False, unicode=True, oid=hstore)
+            except psycopg2.ProgrammingError, e:
+                pass
+
+        if not cached:
+            dsn['cache'] = cache
 
     def close(self):
         self.driver.close()
@@ -151,6 +180,14 @@ class connect(object):
         with self.execute(statement, 1, context) as cursor:
             return cursor.rowcount
 
+    def gen(self, statement):
+        with self.execute(statement, 1) as cursor:
+            while True:
+                fetch = cursor.fetchone()
+                if fetch == None:
+                    break
+                yield fetch
+
     @contextmanager
     def set(self, statement):
         with self.execute(statement, 2) as cursor: