fenton / fenton / db.py

# forward-compat boilerplate
from __future__ import absolute_import
from __future__ import with_statement
__metaclass__ = type

import os
import sqlalchemy as sql
from sqlalchemy import orm

METADATA = sql.MetaData()
ENV_KEY = 'FENTON_UPGRADING'

SCHEMA_HISTORY =  sql.Table(
    'SCHEMA_HISTORY',
    METADATA,
    sql.Column('version', sql.Integer(), primary_key=True, nullable=False),
    sql.Column('author',  sql.String(), nullable=False),
    sql.Column('stamp',   sql.DateTime(timezone=False), nullable=False),
)


def initialize(app):
    os.putenv('PGTZ', 'UTC')
    name = app.config['fenton.db'] + '.db.'
    engine = sql.engine_from_config(app.config, name)
    if not upgrading():
        c = engine.connect()
        version = get_version(c)
        c.close()
        if find_upgrade(app, version):
            from fenton import logging
            logging.log.error('upgrade pending')
    return orm.sessionmaker(bind=engine, autoflush=False)


# command
def upgrade_pending(cx):
    version = get_version(cx.db)
    return find_upgrade(cx.app, version)



# command
def upgrade_offline(cx):
    version = get_version(cx.db)
    up = find_upgrade(cx.app, version)
    return up and not up.online or False



# command
def upgrade(cx, commit=False):
    cx = UpgradeContext(cx)
    version = get_version(cx.db)
    upgrade = find_upgrade(cx.app, version)
    if not upgrade:
        print 'Nothing to do'
        return 1
    try:
        with cx:
            do_upgrade(cx, upgrade)
            print 'Upgraded to version', upgrade.version
            if not commit:
                raise NoCommit
    except NoCommit:
        print 'ROLLBACK'
    else:
        print 'COMMIT'


# decorator
def post_upgrade(f):
    UpgradeContext.post_upgraders.append(f)
    return f


def upgrading(set=None):
    if set:
        os.environ[ENV_KEY] = 'true'
    else:
        return bool(os.environ.get(ENV_KEY))


def get_version(conn):
    v = SCHEMA_HISTORY.c.version
    q = sql.select([v]).order_by(v.desc()).limit(1)
    return conn.execute(q).scalar() or 0


def get_script(v, p):
    ext = os.path.splitext(p)[1][1:].lower()
    C = ext == 'sql' and SqlScript or PyScript
    return C(v, p)


def do_upgrade(cx, upgrade):
    for f in [upgrade] + cx.post_upgraders:
        print f
        f(cx)
        cx.flush()

    q = SCHEMA_HISTORY.insert().values(version=upgrade.version,
                                       author=cx.user.username)
    cx.execute(q)
    cx.flush()


def grant(cx, privs, user):
    run = cx.bind.execute
    objs = run('''
        SELECT relname
        FROM pg_catalog.pg_class c,
             pg_catalog.pg_namespace n
        WHERE c.relkind IN ('r', 'v', 'S')
            AND n.oid = c.relnamespace
            AND n.nspname = 'public' ''')
    objs = [row[0] for row in objs]
    for obj in objs:
        run('GRANT %s ON "%s" TO "%s" ' % (privs, obj, user))


def find_upgrade(app, version, _memo=[]):

    if _memo:
        return _memo[1]

    _memo.append(None)
    OP = os.path
    version += 1
    upgrades = {}
    pending = None
    path = app.config['upgrade.path']

    if not OP.exists(path):
        print 'Directory %s does not exist' % path

    elif not OP.isdir(path):
        print 'Not a directory:', path

    for d, _, files in os.walk(path, followlinks=True):
        for n in files:
            p = OP.join(d, n)
            if not OP.isfile(p):
                continue
            try:
                v = int(OP.splitext(n)[0])
            except ValueError:
                continue
            if v < version:
                continue
            if v in upgrades:
                print 'Duplicate upgrade:', p
                continue
            upgrades[v] = p

    try:
        pending = upgrades.pop(version)
    except KeyError:
        pass
    else:
        pending = get_script(version, pending)
    if upgrades:
        print 'Skipping unexpected upgrades:'
        for v, p in sorted(upgrades.iteritems()):
            print p

    _memo.append(pending)
    return _memo[1]


class NoCommit(Exception):
    'Marker for non-committed upgrade'


class Upgrade:
    def __init__(self, version, path):
        self.version = version
        self.path = path
        self.init()
    def __call__(self, cx):
        return self.execute(cx)
    def __repr__(self):
        return '<Version %d:%s>' % (self.version, self.path)
    def init(self):
        pass


class PyScript(Upgrade):
    online = property(lambda x:x.module.online)
    def init(self):
        src = open(self.path, 'rb').read()
        code = compile(src, self.path, 'exec')
        import imp
        self.module = mod = imp.new_module('version-%d' % self.version)
        self.module.online = False
        mod.__file__ = self.path
        exec code in mod.__dict__

    def execute(self, cx):
        self.module.upgrade(cx)


class SqlScript(Upgrade):
    def execute(self, cx):
        sql = open(self.filename).read()
        cx.run(sql.replace('%', '%%'))


class UpgradeContext:
    post_upgraders = []
    def __init__(self, cx):
        self.__cx = cx
        name = cx.app.config['upgrade.db'] + '.db.'
        engine = sql.engine_from_config(cx.app.config, name)
        cx.app.db.configure(bind=engine, autoflush=True)
        self.context = cx
        self.bind = cx.bind = cx.db.connection()
        self.db = self.bind.db = cx.db
        self.flush = cx.db.flush

    def __getattr__(self, name):
        return getattr(self.__cx, name)

    def __enter__(self):
        return self.__cx.__enter__()

    def __exit__(self, *errors):
        return self.__cx.__exit__(*errors)

    def execute(self, q, **kw):
        self.flush()
        return self.bind.execute(q, **kw)

    def runmany(self, *qq, **kw):
        r = None
        for q in qq:
            if isinstance(q, basestring):
                print q % (kw or {})
                print
            r = self.execute(q, **kw)
        return r

    def run(self, qq, **kw):
        if isinstance(qq, basestring):
            qq = [q for q in qq.split(';\n') if q.strip()]
        elif not isinstance(qq, list):
            qq = [qq]
        return self.runmany(*qq, **kw)

    def create(self, *classes):
        ts = [c.__table__ for c in classes]
        METADATA.create_all(bind=self.bind, tables=ts or None)

    def drop_columns(self, *colnames):
        if len(colnames) == 1 and '\n' in colnames[0]:
            colnames = colnames[0].split()
        for n in colnames:
            self.run('ALTER TABLE %s DROP %s CASCADE' % tuple(n.split('.', 1)))

    def drop_tables(self, *tablenames):
        if len(tablenames) == 1 and '\n' in tablenames[0]:
            tablenames = tablenames[0].split()
        for t in tablenames:
            self.run('DROP TABLE %s CASCADE' % t)
            self.delete_metainfo(t)

    def insert_metainfo(self, *classes):
        from fenton import data
        for c in classes:
            data._insert_metainfo(self, c)

    def delete_metainfo(self, *tablenames):
        from fenton import data
        C = data.MetaObjectClass
        q = C.__table__.delete()
        for t in tablenames:
            self.run(q.where(C.tablename == t))
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.