__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 class connection(object):
  16     def __init__(self, driver):
  17         self.driver = driver
  18
  19     @contextmanager
  20     def cursor(self):
  21         try:
  22             cursor = self.driver.cursor(cursor_factory=psycopg2.extras.DictCursor)
  23             yield cursor
  24         finally:
  25             cursor.close()
  26
  27     @contextmanager
  28     def execute(self, statement, depth=0):
  29         with self.cursor() as cursor:
  30             # two frames, accounting for execute() and @contextmanager
  31             locals = inspect.currentframe(depth + 2).f_locals
  32             try:
  33                 cursor.execute(statement.format(**locals), locals)
  34             finally:
  35                 del locals
  36             yield cursor
  37
  38     @contextmanager
  39     def transact(self, synchronous_commit=True):
  40         with self.cursor() as cursor:
  41             if not synchronous_commit:
  42                 cursor.execute('set local synchronous_commit = off')
  43
  44         try:
  45             yield transaction(self)
  46             self.driver.commit()
  47         except:
  48             self.driver.rollback()
  49             raise
  50
  51 class transaction(object):
  52     def __init__(self, connection):
  53         self.connection = connection
  54
  55     def pull(self, statement):
  56         with self.connection.execute(statement, 1) as cursor:
  57             return cursor.fetchall()
  58
  59     def yank(self, statement, offset=0):
  60         with self.connection.execute(statement, 1 + offset) as cursor:
  61             rows = cursor.fetchall()
  62             return rows[0] if len(rows) != 0 else None
  63
  64     def push(self, statement):
  65         with self.connection.execute(statement, 1) as cursor:
  66             pass
  67
  68     def exists(self, statement):
  69         return self.yank('''
  70             select exists (
  71                 {statement}
  72             )
  73         '''.format(**locals()), 1)[0]
  74
  75 @contextmanager
  76 def connect(dsn):
  77     attempt = 0
  78     while True:
  79         try:
  80             driver = psycopg2.connect(**dsn)
  81             break
  82         except psycopg2.OperationalError, e:
  83             if attempt == 2:
  84                 raise e
  85             attempt = attempt + 1
  86
  87     try:
  88         driver.set_client_encoding('UNICODE')
  89         yield connection(driver)
  90     finally:
  91         driver.close()
  92
  93 def connected(dsn):
  94     def wrapped(method):
  95         def replaced(*args, **kw):
  96             with connect(dsn) as connection:
  97                 return method(connection, *args, **kw)
  98         return replaced
  99     return wrapped
 100
 101 @contextmanager
 102 def transact(dsn, **args):
 103     with connect(dsn) as connection:
 104         with connection.transact(**args) as cursor:
 105             yield cursor
 106
 107 """
 108 def slap_(sql, table, keys, values, path):
 109     csr = sql.cursor()
 110     try:
 111         csr.execute('savepoint iou')
 112         try:
 113             both = dict(keys, **values)
 114             fields = both.keys()
 115
 116             csr.execute('''
 117                 insert into %s (%s) values (%s)
 118             ''' % (
 119                 table,
 120                 ', '.join(fields),
 121                 ', '.join(['%s' for key in fields])
 122             ), both.values())
 123         except psycopg2.IntegrityError, e:
 124             csr.execute('rollback to savepoint iou')
 125
 126             csr.execute('''
 127                 update %s set %s where %s
 128             ''' % (
 129                 table,
 130                 ', '.join([
 131                     key + ' = %s'
 132                 for key in values.keys()]),
 133                 ' and '.join([
 134                     key + ' = %s'
 135                 for key in keys.keys()])
 136             ), values.values() + keys.values())
 137
 138         return path_(csr, path)
 139     finally:
 140         csr.close()
 141 """