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