__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             cursor.execute(statement.format(**locals), locals)
  33             yield cursor
  34
  35     @contextmanager
  36     def transact(self, synchronous_commit=True):
  37         with self.cursor() as cursor:
  38             if not synchronous_commit:
  39                 cursor.execute('set local synchronous_commit = off')
  40
  41         try:
  42             yield transaction(self)
  43             self.driver.commit()
  44         except:
  45             self.driver.rollback()
  46             raise
  47
  48 class transaction(object):
  49     def __init__(self, connection):
  50         self.connection = connection
  51
  52     def pull(self, statement):
  53         with self.connection.execute(statement, 1) as cursor:
  54             return cursor.fetchall()
  55
  56     def yank(self, statement):
  57         with self.connection.execute(statement, 1) as cursor:
  58             rows = cursor.fetchall()
  59             return rows[0] if len(rows) != 0 else None
  60
  61     def push(self, statement):
  62         with self.connection.execute(statement, 1) as cursor:
  63             pass
  64
  65 @contextmanager
  66 def connect(dsn):
  67     attempt = 0
  68     while True:
  69         try:
  70             driver = psycopg2.connect(**dsn)
  71             break
  72         except psycopg2.OperationalError, e:
  73             if attempt == 2:
  74                 raise e
  75             attempt = attempt + 1
  76
  77     try:
  78         driver.set_client_encoding('UNICODE')
  79         yield connection(driver)
  80     finally:
  81         driver.close()
  82
  83 @contextmanager
  84 def transact(dsn, **args):
  85     with connect(dsn) as connection:
  86         with connection.transact(**args) as cursor:
  87             yield cursor
  88
  89 """
  90 def slap_(sql, table, keys, values, path):
  91     csr = sql.cursor()
  92     try:
  93         csr.execute('savepoint iou')
  94         try:
  95             both = dict(keys, **values)
  96             fields = both.keys()
  97
  98             csr.execute('''
  99                 insert into %s (%s) values (%s)
 100             ''' % (
 101                 table,
 102                 ', '.join(fields),
 103                 ', '.join(['%s' for key in fields])
 104             ), both.values())
 105         except psycopg2.IntegrityError, e:
 106             csr.execute('rollback to savepoint iou')
 107
 108             csr.execute('''
 109                 update %s set %s where %s
 110             ''' % (
 111                 table,
 112                 ', '.join([
 113                     key + ' = %s'
 114                 for key in values.keys()]),
 115                 ' and '.join([
 116                     key + ' = %s'
 117                 for key in keys.keys()])
 118             ), values.values() + keys.values())
 119
 120         return path_(csr, path)
 121     finally:
 122         csr.close()
 123 """