This interface is actually getting closer to what I want.

[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)

def one(values):
    if values == None or len(values) == 0:
        return None
    else:
        assert len(values) == 1
        return values[0]

class connection(object):
    def __init__(self, driver):
        self.driver = driver
        self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)

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

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

    @contextmanager
    def transact(self, synchronous_commit=True):
        self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
        try:
            with self.cursor() as cursor:
                if not synchronous_commit:
                    cursor.execute('set local synchronous_commit = off')

            yield
            self.driver.commit()
        except:
            self.driver.rollback()
            raise
        finally:
            self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)

    def one_(self, statement):
        with self.execute(statement, 2) as cursor:
            one = cursor.fetchone()
            if one == None:
                return None

            assert cursor.fetchone() == None
            return one

    def __call__(self, procedure, *parameters):
        with self.execute(statement, 1) as cursor:
            return cursor.callproc(procedure, *parameters)

    def run(self, statement, locals=None):
        with self.execute(statement, 1, locals) as cursor:
            return cursor.rowcount

    @contextmanager
    def set(self, statement):
        with self.execute(statement, 1) as cursor:
            yield cursor

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

    def one(self, statement):
        return self.one_(statement)

    def has(self, statement):
        return one(self.one_('select exists(%s)' % (statement,)))

@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):
            yield connection

"""
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()
"""