sqlhelpers.py

from sqlalchemy.sql.expression import Insert
from sqlalchemy.ext.compiler import compiles


class Upsert(Insert):
    '''Compiles an upsert expression for use with mysql.

    Requires a keyword arg ``upsert_fields`` which lists field names to be updated in
    the ON DUPLICATE KEY UPDATE clause.
    '''

    def __init__(self, *args, **kwargs):
        self.upsert_fields = kwargs.pop('upsert_fields')
        assert self.upsert_fields, ('Must specify at least one field to be '
                                    'updated in case of duplicate entry.')
        super(Upsert, self).__init__(*args, **kwargs)


@compiles(Upsert, 'mysql')
def visit_upsert(element, compiler, **kwargs):
    dupe_values = ', '.join('{field} = VALUES(`{field}`)'.format(field=field)
                            for field in element.upsert_fields)
    return '{insert_clause} ON DUPLICATE KEY UPDATE {dupe_values}'.format(
        insert_clause=compiler.visit_insert(element, **kwargs),
        dupe_values=dupe_values
    )