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