two_phase_commit.py

# -*- coding: utf-8 -*-

from johnny.middleware import QueryCacheMiddleware
from johnny.transaction import TransactionManager
from johnny.cache import QueryCacheBackend
from uuid import uuid4
from django.db import connections, DEFAULT_DB_ALIAS, load_backend

# not correctness-critical, after expiry will force complete cache flush
XACT_TIMEOUT = 365 * 86400

def _cache_backend():
    # borg pattern, all instances share the same state
    return QueryCacheMiddleware().query_cache_backend

def make_xact_id(xact_id):
    return 'XACT_{0}'.format(xact_id)

def prepare_transaction(using=None):
    backend = _cache_backend()
    assert isinstance(backend, QueryCacheBackend)
    backend = backend.cache_backend

    xact_id = uuid4()
    effective_xact_id = make_xact_id(xact_id)
    if using is None:
        _using = DEFAULT_DB_ALIAS
    else:
        _using = using
    conn = connections[_using]
    cur = conn.cursor()
    cur.execute('PREPARE TRANSACTION %s', (effective_xact_id,))

    # fold all savepoint data into the local cache
    backend._commit_all_savepoints()


    c = backend.local.get('invalidated_{0}'.format(_using))
    backend.cache_backend.set(effective_xact_id, c, XACT_TIMEOUT)
    backend._clear(using)
    backend._clear_sid_stack(using)
    return xact_id

def clean_connection(conn_alias=None):
    if conn_alias is None:
        conn_alias = DEFAULT_DB_ALIAS
    db_setup = dict(connections.databases[conn_alias]) # shallow copy
    db_setup['OPTIONS'] = dict(autocommit=True)
    db_backend = load_backend(db_setup['ENGINE'])
    return db_backend.DatabaseWrapper(db_setup, '{0}__autocommit'.format(conn_alias))

def commit_prepared(xact_id, using=None):
    conn = clean_connection(using)
    backend = _cache_backend()
    assert isinstance(backend, QueryCacheBackend)

    effective_xact_id = make_xact_id(xact_id)
    cur = conn.cursor()
    cur.execute('COMMIT PREPARED %s', (effective_xact_id,))

    to_invalidate = backend.cache_backend.cache_backend.get(effective_xact_id)
    if to_invalidate is None:
        # oops. invalidate everything, just in case
        backend.flush_query_cache()
    else:
        for table in to_invalidate:
            backend.keyhandler.invalidate_table(table)

def rollback_prepared(xact_id, using=None):
    conn = clean_connection(using)
    backend = _cache_backend()
    assert isinstance(backend, QueryCacheBackend)

    effective_xact_id = make_xact_id(xact_id)
    cur = conn.cursor()
    cur.execute('ROLLBACK PREPARED %s', (effective_xact_id,))

    backend.cache_backend.cache_backend.delete(effective_xact_id)