Add execute_() and push_() that take a manual locals.

[cyql.git] / __init__.py

from __future__ import unicode_literals
from __future__ import print_function

import inspect
import os

from contextlib import contextmanager

import psycopg2
import psycopg2.extras
import psycopg2.pool

psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)

class connection(object):
    def __init__(self, driver):
        self.driver = driver

    @contextmanager
    def cursor(self):
        try:
            cursor = self.driver.cursor(cursor_factory=psycopg2.extras.DictCursor)
            yield cursor
        finally:
            cursor.close()

    @contextmanager
    def execute_(self, statement, locals):
        with self.cursor() as cursor:
            cursor.execute(statement.format(**locals), locals)
            yield cursor

    @contextmanager
    def execute(self, statement, depth=0):
        with self.cursor() as cursor:
            # two frames, accounting for execute() and @contextmanager
            locals = inspect.currentframe(depth + 2).f_locals
            try:
                cursor.execute(statement.format(**locals), locals)
            finally:
                del locals
            yield cursor

    @contextmanager
    def transact(self, synchronous_commit=True):
        with self.cursor() as cursor:
            if not synchronous_commit:
                cursor.execute('set local synchronous_commit = off')

        try:
            yield transaction(self)
            self.driver.commit()
        except:
            self.driver.rollback()
            raise

class transaction(object):
    def __init__(self, connection):
        self.connection = connection

    def pull(self, statement):
        with self.connection.execute(statement, 1) as cursor:
            return cursor.fetchall()

    def yank(self, statement, offset=0):
        with self.connection.execute(statement, 1 + offset) as cursor:
            rows = cursor.fetchall()
            return rows[0] if len(rows) != 0 else None

    def push(self, statement):
        with self.connection.execute(statement, 1) as cursor:
            pass

    def push_(self, statement, locals):
        with self.connection.execute_(statement, locals) as cursor:
            pass

    def exists(self, statement):
        return self.yank('''
            select exists (
                {statement}
            )
        '''.format(**locals()), 1)[0]

@contextmanager
def connect(dsn):
    attempt = 0
    while True:
        try:
            driver = psycopg2.connect(**dsn)
            break
        except psycopg2.OperationalError, e:
            if attempt == 2:
                raise e
            attempt = attempt + 1

    try:
        driver.set_client_encoding('UNICODE')
        yield connection(driver)
    finally:
        driver.close()

def connected(dsn):
    def wrapped(method):
        def replaced(*args, **kw):
            with connect(dsn) as connection:
                return method(connection, *args, **kw)
        return replaced
    return wrapped

@contextmanager
def transact(dsn, **args):
    with connect(dsn) as connection:
        with connection.transact(**args) as cursor:
            yield cursor

"""
def slap_(sql, table, keys, values, path):
    csr = sql.cursor()
    try:
        csr.execute('savepoint iou')
        try:
            both = dict(keys, **values)
            fields = both.keys()

            csr.execute('''
                insert into %s (%s) values (%s)
            ''' % (
                table,
                ', '.join(fields),
                ', '.join(['%s' for key in fields])
            ), both.values())
        except psycopg2.IntegrityError, e:
            csr.execute('rollback to savepoint iou')

            csr.execute('''
                update %s set %s where %s
            ''' % (
                table,
                ', '.join([
                    key + ' = %s'
                for key in values.keys()]),
                ' and '.join([
                    key + ' = %s'
                for key in keys.keys()])
            ), values.values() + keys.values())

        return path_(csr, path)
    finally:
        csr.close()
"""