
"""Wrapper around psycopg1/2.

Preferred is psycopg2, fallback to psycopg1.

Interface provided is psycopg1:
    - dict* methods.
    - new columns can be assigned to row.

"""

import sys

__all__ = []

try:
    ##from psycopg2.psycopg1 import connect as _pgconnect
    # psycopg2.psycopg1.cursor is too backwards compatible,
    # to the point of avoiding optimized access.
    # only backwards compat thing we need is dict* methods

    import psycopg2.extensions, psycopg2.extras
    from sqltools import dbdict

    class _CompatRow(psycopg2.extras.DictRow):
        """Make DictRow more dict-like."""
        __slots__ = ('_index',)

        def __contains__(self, k):
            """Returns if such row has such column."""
            return k in self._index

        def copy(self):
            """Return regular dict."""
            return dbdict(self.iteritems())
        
        def iterkeys(self):
            return self._index.iterkeys()

        def itervalues(self):
            return list.__iter__(self)

        # obj.foo access
        def __getattr__(self, k):
            return self[k]

    class _CompatCursor(psycopg2.extras.DictCursor):
        """Regular psycopg2 DictCursor with dict* methods."""
        def __init__(self, *args, **kwargs):
            psycopg2.extras.DictCursor.__init__(self, *args, **kwargs)
            self.row_factory = _CompatRow
        dictfetchone = psycopg2.extras.DictCursor.fetchone
        dictfetchall = psycopg2.extras.DictCursor.fetchall
        dictfetchmany = psycopg2.extras.DictCursor.fetchmany
    
    class _CompatConnection(psycopg2.extensions.connection):
        """Connection object that uses _CompatCursor."""
        def cursor(self):
            return psycopg2.extensions.connection.cursor(self, cursor_factory = _CompatCursor)
    
    def _pgconnect(cstr):
        """Create a psycopg2 connection."""
        return _CompatConnection(cstr)

except ImportError:
    # use psycopg 1
    try:
        from psycopg import connect as _pgconnect
    except ImportError:
        print "Please install psycopg2 module"
        sys.exit(1)

def connect_database(connstr):
    """Create a db connection with connect_timeout option.
    
    Default connect_timeout is 15, to change put it directly into dsn.
    """

    # allow override
    if connstr.find("connect_timeout") < 0:
        connstr += " connect_timeout=15"

    # create connection
    return _pgconnect(connstr)


@@ -198,25 +198,16 @@
 _ = %5B%5D%0A%0A
-try:%0A    
 ##from p
@@ -252,20 +252,16 @@
 connect%0A
-    
 # psycop
@@ -308,20 +308,16 @@
 atible,%0A
-    
 # to the
@@ -353,20 +353,16 @@
 access.%0A
-    
 # only b
@@ -409,20 +409,16 @@
 ethods%0A%0A
-    
 import p
@@ -453,25 +453,30 @@
 .extras%0A
-    
 from 
+skytools.
 sqltools
@@ -491,20 +491,16 @@
 dbdict%0A%0A
-    
 class _C
@@ -538,20 +538,16 @@
 w):%0A    
-  
-  
 %22%22%22Make 
@@ -569,28 +569,24 @@
 ct-like.%22%22%22%0A
-    
     __slots_
@@ -602,28 +602,24 @@
 dex',)%0A%0A    
- 
-   
 def __contai
@@ -625,36 +625,32 @@
 ins__(self, k):%0A
-    
         %22%22%22Retur
@@ -688,28 +688,24 @@
 %22%22%22%0A        
-    
 return k in 
@@ -717,20 +717,16 @@
 _index%0A%0A
-    
     def 
@@ -741,28 +741,24 @@
 f):%0A        
-    
 %22%22%22Return re
@@ -776,28 +776,24 @@
 %22%22%22%0A        
-   
- 
 return dbdic
@@ -820,21 +820,13 @@
     
-    %0A  
+%0A
     
-  
 def 
@@ -833,36 +833,32 @@
 iterkeys(self):%0A
-    
         return s
@@ -880,28 +880,24 @@
 keys()%0A%0A    
- 
-   
 def itervalu
@@ -898,36 +898,32 @@
 ervalues(self):%0A
-    
         return l
@@ -946,20 +946,16 @@
 f)%0A%0A    
- 
-   
 # obj.fo
@@ -963,20 +963,16 @@
  access%0A
-    
     def 
@@ -997,28 +997,24 @@
 k):%0A        
-    
 return self%5B
@@ -1017,20 +1017,16 @@
 elf%5Bk%5D%0A%0A
-    
 class _C
@@ -1062,28 +1062,24 @@
 ictCursor):%0A
-    
     %22%22%22Regul
@@ -1124,28 +1124,24 @@
 ods.%22%22%22%0A    
-  
-  
 def __init__
@@ -1165,20 +1165,16 @@
 wargs):%0A
-    
         
@@ -1240,20 +1240,16 @@
         
-  
-  
 self.row
@@ -1270,20 +1270,16 @@
 mpatRow%0A
-    
     dict
@@ -1325,28 +1325,24 @@
 etchone%0A    
- 
-   
 dictfetchall
@@ -1380,20 +1380,16 @@
 etchall%0A
-    
     dict
@@ -1437,25 +1437,17 @@
 tchmany%0A
-    %0A    
+%0A
 class _C
@@ -1495,20 +1495,16 @@
 ction):%0A
-    
     %22%22%22C
@@ -1552,20 +1552,16 @@
 %22%22%22%0A    
-   
- 
 def curs
@@ -1566,28 +1566,24 @@
 rsor(self):%0A
-    
         retu
@@ -1664,313 +1664,8 @@
 sor)
-%0A    %0A    def _pgconnect(cstr):%0A        %22%22%22Create a psycopg2 connection.%22%22%22%0A        return _CompatConnection(cstr)%0A%0Aexcept ImportError:%0A    # use psycopg 1%0A    try:%0A        from psycopg import connect as _pgconnect%0A    except ImportError:%0A        print %22Please install psycopg2 module%22%0A        sys.exit(1)
 %0A%0Ade
@@ -1980,25 +1980,32 @@
 return _
-pgc
+CompatC
 onnect
+ion
 (connstr

drop support for psycopg1