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