__init__.py

   1 from __future__ import absolute_import
   2 from __future__ import division
   3 from __future__ import print_function
   4 from __future__ import unicode_literals
   5
   6 from future_builtins import ascii, filter, hex, map, oct, zip
   7
   8 import inspect
   9 import os
  10
  11 from contextlib import contextmanager
  12
  13 import psycopg2
  14 import psycopg2.extras
  15 import psycopg2.pool
  16
  17 psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
  18 psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
  19
  20 class connect(object):
  21     def __init__(self, dsn):
  22         attempt = 0
  23         while True:
  24             try:
  25                 self.driver = psycopg2.connect(**dsn)
  26                 break
  27             except psycopg2.OperationalError, e:
  28                 if attempt == 2:
  29                     raise e
  30                 attempt = attempt + 1
  31
  32         try:
  33             self.driver.set_client_encoding('UNICODE')
  34             self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
  35         except:
  36             self.driver.close()
  37
  38         try:
  39             psycopg2.extras.register_hstore(self.driver, globally=False, unicode=True)
  40         except ProgrammingError, e:
  41             pass
  42
  43     def close(self):
  44         self.driver.close()
  45
  46     def __enter__(self):
  47         return self
  48
  49     def __exit__(self, type, value, traceback):
  50         self.close()
  51
  52     @contextmanager
  53     def cursor(self):
  54         cursor = self.driver.cursor(cursor_factory=psycopg2.extras.DictCursor)
  55         try:
  56             yield cursor
  57         finally:
  58             cursor.close()
  59
  60     @contextmanager
  61     def execute(self, statement, depth=0, context=None):
  62         # two frames, accounting for execute() and @contextmanager
  63         frame = inspect.currentframe(depth + 2)
  64
  65         with self.cursor() as cursor:
  66             f_globals = None
  67             f_locals = frame.f_locals
  68
  69             if context == None:
  70                 context = dict(**f_locals)
  71
  72             start = 0
  73             while True:
  74                 percent = statement.find('%', start)
  75                 if percent == -1:
  76                     break
  77
  78                 next = statement[percent + 1]
  79                 if next == '(':
  80                     start = statement.index(')', percent + 2) + 2
  81                     assert statement[start - 1] == 's'
  82                 elif next == '{':
  83                     start = statement.index('}', percent + 2)
  84                     assert statement[start + 1] == 's'
  85                     code = statement[percent + 2:start]
  86
  87                     if f_globals == None:
  88                         f_globals = frame.f_globals
  89
  90                     key = '__cyql__%i' % (percent,)
  91                     # XXX: compile() in the frame's context
  92                     context[key] = eval(code, f_globals, f_locals)
  93
  94                     statement = '%s%%(%s)%s' % (statement[0:percent], key, statement[start + 1:])
  95                     start = percent + len(key) + 4
  96                 elif next in ('%', 's'):
  97                     start = percent + 2
  98                 else:
  99                     assert False
 100
 101             cursor.execute(statement, context)
 102
 103             del context
 104             del f_locals
 105             del f_globals
 106
 107             yield cursor
 108
 109     @contextmanager
 110     def transact(self, synchronous_commit=True):
 111         self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
 112         try:
 113             with self.cursor() as cursor:
 114                 if not synchronous_commit:
 115                     cursor.execute('set local synchronous_commit = off')
 116
 117             yield
 118             self.driver.commit()
 119         except:
 120             self.driver.rollback()
 121             raise
 122         finally:
 123             self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
 124
 125     def one_(self, statement, context=None):
 126         with self.execute(statement, 2, context) as cursor:
 127             one = cursor.fetchone()
 128             if one == None:
 129                 return None
 130
 131             assert cursor.fetchone() == None
 132             return one
 133
 134     def __call__(self, procedure, *parameters):
 135         with self.execute(statement, 1) as cursor:
 136             return cursor.callproc(procedure, *parameters)
 137
 138     def run(self, statement, context=None):
 139         with self.execute(statement, 1, context) as cursor:
 140             return cursor.rowcount
 141
 142     @contextmanager
 143     def set(self, statement):
 144         with self.execute(statement, 1) as cursor:
 145             yield cursor
 146
 147     def all(self, statement):
 148         with self.execute(statement, 1) as cursor:
 149             return cursor.fetchall()
 150
 151     def one(self, statement, context=None):
 152         return self.one_(statement, context)
 153
 154     def has(self, statement):
 155         exists, = self.one_('select exists(%s)' % (statement,))
 156         return exists
 157
 158 def connected(dsn):
 159     def wrapped(method):
 160         def replaced(*args, **kw):
 161             with connect(dsn) as sql:
 162                 return method(*args, sql=sql, **kw)
 163         return replaced
 164     return wrapped
 165
 166 @contextmanager
 167 def transact(dsn, *args, **kw):
 168     with connect(dsn) as connection:
 169         with connection.transact(*args, **kw):
 170             yield connection
 171
 172 """
 173 def slap_(sql, table, keys, values, path):
 174     csr = sql.cursor()
 175     try:
 176         csr.execute('savepoint iou')
 177         try:
 178             both = dict(keys, **values)
 179             fields = both.keys()
 180
 181             csr.execute('''
 182                 insert into %s (%s) values (%s)
 183             ''' % (
 184                 table,
 185                 ', '.join(fields),
 186                 ', '.join(['%s' for key in fields])
 187             ), both.values())
 188         except psycopg2.IntegrityError, e:
 189             csr.execute('rollback to savepoint iou')
 190
 191             csr.execute('''
 192                 update %s set %s where %s
 193             ''' % (
 194                 table,
 195                 ', '.join([
 196                     key + ' = %s'
 197                 for key in values.keys()]),
 198                 ' and '.join([
 199                     key + ' = %s'
 200                 for key in keys.keys()])
 201             ), values.values() + keys.values())
 202
 203         return path_(csr, path)
 204     finally:
 205         csr.close()
 206 """