__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         with self.cursor() as cursor:
  57             # two frames, accounting for execute() and @contextmanager
  58             locals = inspect.currentframe(depth + 2).f_locals
  59             try:
  60                 if context == None:
  61                     context = locals
  62                 cursor.execute(statement.format(**locals), context)
  63             finally:
  64                 del locals
  65             yield cursor
  66
  67     @contextmanager
  68     def transact(self, synchronous_commit=True):
  69         self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
  70         try:
  71             with self.cursor() as cursor:
  72                 if not synchronous_commit:
  73                     cursor.execute('set local synchronous_commit = off')
  74
  75             yield
  76             self.driver.commit()
  77         except:
  78             self.driver.rollback()
  79             raise
  80         finally:
  81             self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
  82
  83     def one_(self, statement):
  84         with self.execute(statement, 2) as cursor:
  85             one = cursor.fetchone()
  86             if one == None:
  87                 return None
  88
  89             assert cursor.fetchone() == None
  90             return one
  91
  92     def __call__(self, procedure, *parameters):
  93         with self.execute(statement, 1) as cursor:
  94             return cursor.callproc(procedure, *parameters)
  95
  96     def run(self, statement, locals=None):
  97         with self.execute(statement, 1, locals) as cursor:
  98             return cursor.rowcount
  99
 100     @contextmanager
 101     def set(self, statement):
 102         with self.execute(statement, 1) as cursor:
 103             yield cursor
 104
 105     def all(self, statement):
 106         with self.execute(statement, 1) as cursor:
 107             return cursor.fetchall()
 108
 109     def one(self, statement):
 110         return self.one_(statement)
 111
 112     def has(self, statement):
 113         exists, = self.one_('select exists(%s)' % (statement,))
 114         return exists
 115
 116 def connected(dsn):
 117     def wrapped(method):
 118         def replaced(*args, **kw):
 119             with connect(dsn) as sql:
 120                 return method(*args, sql=sql, **kw)
 121         return replaced
 122     return wrapped
 123
 124 @contextmanager
 125 def transact(dsn, *args, **kw):
 126     with connect(dsn) as connection:
 127         with connection.transact(*args, **kw):
 128             yield connection
 129
 130 """
 131 def slap_(sql, table, keys, values, path):
 132     csr = sql.cursor()
 133     try:
 134         csr.execute('savepoint iou')
 135         try:
 136             both = dict(keys, **values)
 137             fields = both.keys()
 138
 139             csr.execute('''
 140                 insert into %s (%s) values (%s)
 141             ''' % (
 142                 table,
 143                 ', '.join(fields),
 144                 ', '.join(['%s' for key in fields])
 145             ), both.values())
 146         except psycopg2.IntegrityError, e:
 147             csr.execute('rollback to savepoint iou')
 148
 149             csr.execute('''
 150                 update %s set %s where %s
 151             ''' % (
 152                 table,
 153                 ', '.join([
 154                     key + ' = %s'
 155                 for key in values.keys()]),
 156                 ' and '.join([
 157                     key + ' = %s'
 158                 for key in keys.keys()])
 159             ), values.values() + keys.values())
 160
 161         return path_(csr, path)
 162     finally:
 163         csr.close()
 164 """