from . import Base
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager
from sqlalchemy import create_engine
from models import * # noqa


class SQLAlchemyDriver(object):
    def __init__(self, conn, **config):
        self.engine = create_engine(conn, **config)

        Base.metadata.bind = self.engine
        Base.metadata.create_all()

        self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)

    @property
    @contextmanager
    def session(self):
        '''
        Provide a transactional scope around a series of operations.
        '''
        session = self.Session()
        yield session

        try:
            session.commit()
        except:
            session.rollback()
            raise
        finally:
            session.close()

    def get_or_create(self, session, model, query, props=None):
        '''
        Get or create a row
        Args:
            session: sqlalchemy session
            model: the ORM class from userdatamodel.models
            query: a dict of query parameters
            props: extra props aside from query to be added to the object on
                   creation
        Returns:
            result object of the model class
        '''
        result = session.query(model).filter_by(**query).first()
        if result is None:
            args = props if props is not None else {}
            args.update(query)
            result = model(**args)
            session.add(result)
        return result

@@ -327,44 +327,8 @@
 ine%0A
-        Base.metadata.create_all()%0A%0A
     
@@ -401,16 +401,79 @@
 =False)%0A
+        self.pre_migrate()%0A        Base.metadata.create_all()%0A%0A
 %0A    @pr
@@ -1532,8 +1532,579 @@
  result%0A
+%0A    def pre_migrate(self):%0A        '''%0A        migration script to be run before create_all%0A        '''%0A        if not self.engine.dialect.supports_alter:%0A            print(%0A                %22This engine dialect doesn't support altering%22%0A                %22 so we are not migrating even if necessary!%22)%0A            return%0A%0A        if not self.engine.dialect.has_table(self.engine, 'Group'):%0A            print(%22Altering table research_group to group%22)%0A            with self.session as session:%0A                session.execute('ALTER TABLE research_group rename to %22Group%22')%0A

migrate research_group