__init__.py

   1 from __future__ import unicode_literals
   2 from __future__ import print_function
   3
   4 import inspect
   5 import os
   6
   7 from contextlib import contextmanager
   8
   9 import psycopg2
  10 import psycopg2.extras
  11 import psycopg2.pool
  12
  13 psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
  14
  15 def one(values):
  16     if values == None or len(values) == 0:
  17         return None
  18     else:
  19         assert len(values) == 1
  20         return values[0]
  21
  22 class connect(object):
  23     def __init__(self, dsn):
  24         attempt = 0
  25         while True:
  26             try:
  27                 self.driver = psycopg2.connect(**dsn)
  28                 break
  29             except psycopg2.OperationalError, e:
  30                 if attempt == 2:
  31                     raise e
  32                 attempt = attempt + 1
  33
  34         try:
  35             self.driver.set_client_encoding('UNICODE')
  36             self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
  37         except:
  38             self.driver.close()
  39
  40     def close(self):
  41         self.driver.close()
  42
  43     def __enter__(self):
  44         return self
  45
  46     def __exit__(self, type, value, traceback):
  47         self.close()
  48
  49     @contextmanager
  50     def cursor(self):
  51         cursor = self.driver.cursor(cursor_factory=psycopg2.extras.DictCursor)
  52         try:
  53             yield cursor
  54         finally:
  55             cursor.close()
  56
  57     @contextmanager
  58     def execute(self, statement, depth=0, context=None):
  59         with self.cursor() as cursor:
  60             # two frames, accounting for execute() and @contextmanager
  61             locals = inspect.currentframe(depth + 2).f_locals
  62             try:
  63                 if context == None:
  64                     context = locals
  65                 cursor.execute(statement.format(**locals), context)
  66             finally:
  67                 del locals
  68             yield cursor
  69
  70     @contextmanager
  71     def transact(self, synchronous_commit=True):
  72         self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
  73         try:
  74             with self.cursor() as cursor:
  75                 if not synchronous_commit:
  76                     cursor.execute('set local synchronous_commit = off')
  77
  78             yield
  79             self.driver.commit()
  80         except:
  81             self.driver.rollback()
  82             raise
  83         finally:
  84             self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
  85
  86     def one_(self, statement):
  87         with self.execute(statement, 2) as cursor:
  88             one = cursor.fetchone()
  89             if one == None:
  90                 return None
  91
  92             assert cursor.fetchone() == None
  93             return one
  94
  95     def __call__(self, procedure, *parameters):
  96         with self.execute(statement, 1) as cursor:
  97             return cursor.callproc(procedure, *parameters)
  98
  99     def run(self, statement, locals=None):
 100         with self.execute(statement, 1, locals) as cursor:
 101             return cursor.rowcount
 102
 103     @contextmanager
 104     def set(self, statement):
 105         with self.execute(statement, 1) as cursor:
 106             yield cursor
 107
 108     def all(self, statement):
 109         with self.execute(statement, 1) as cursor:
 110             return cursor.fetchall()
 111
 112     def one(self, statement):
 113         return self.one_(statement)
 114
 115     def has(self, statement):
 116         return one(self.one_('select exists(%s)' % (statement,)))
 117
 118 def connected(dsn):
 119     def wrapped(method):
 120         def replaced(*args, **kw):
 121             with connect(dsn) as sql:
 122                 return method(*args, sql=sql, **kw)
 123         return replaced
 124     return wrapped
 125
 126 @contextmanager
 127 def transact(dsn, *args, **kw):
 128     with connect(dsn) as connection:
 129         with connection.transact(*args, **kw):
 130             yield connection
 131
 132 """
 133 def slap_(sql, table, keys, values, path):
 134     csr = sql.cursor()
 135     try:
 136         csr.execute('savepoint iou')
 137         try:
 138             both = dict(keys, **values)
 139             fields = both.keys()
 140
 141             csr.execute('''
 142                 insert into %s (%s) values (%s)
 143             ''' % (
 144                 table,
 145                 ', '.join(fields),
 146                 ', '.join(['%s' for key in fields])
 147             ), both.values())
 148         except psycopg2.IntegrityError, e:
 149             csr.execute('rollback to savepoint iou')
 150
 151             csr.execute('''
 152                 update %s set %s where %s
 153             ''' % (
 154                 table,
 155                 ', '.join([
 156                     key + ' = %s'
 157                 for key in values.keys()]),
 158                 ' and '.join([
 159                     key + ' = %s'
 160                 for key in keys.keys()])
 161             ), values.values() + keys.values())
 162
 163         return path_(csr, path)
 164     finally:
 165         csr.close()
 166 """