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