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