mediagoblin/db/open.py

   1 # GNU MediaGoblin -- federated, autonomous media hosting
   2 # Copyright (C) 2011, 2012 MediaGoblin contributors.  See AUTHORS.
   3 #
   4 # This program is free software: you can redistribute it and/or modify
   5 # it under the terms of the GNU Affero General Public License as published by
   6 # the Free Software Foundation, either version 3 of the License, or
   7 # (at your option) any later version.
   8 #
   9 # This program is distributed in the hope that it will be useful,
  10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 # GNU Affero General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU Affero General Public License
  15 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16
  17
  18 from contextlib import contextmanager
  19 import logging
  20
  21 import six
  22 from sqlalchemy import create_engine, event
  23
  24 from mediagoblin import mg_globals
  25 from mediagoblin.db.base import Base
  26
  27 _log = logging.getLogger(__name__)
  28
  29 from mediagoblin.tools.transition import DISABLE_GLOBALS
  30
  31 def set_models_as_attributes(obj):
  32     """
  33     Set all models as attributes on this object, for convenience
  34
  35     TODO: This should eventually be deprecated.
  36     """
  37     for k, v in six.iteritems(Base._decl_class_registry):
  38         setattr(obj, k, v)
  39
  40
  41 if not DISABLE_GLOBALS:
  42     from mediagoblin.db.base import Session
  43
  44     class DatabaseMaster(object):
  45         def __init__(self, engine):
  46             self.engine = engine
  47
  48             set_models_as_attributes(self)
  49
  50         def commit(self):
  51             Session.commit()
  52
  53         def save(self, obj):
  54             Session.add(obj)
  55             Session.flush()
  56
  57         def check_session_clean(self):
  58             for dummy in Session():
  59                 _log.warn("STRANGE: There are elements in the sql session. "
  60                           "Please report this and help us track this down.")
  61                 break
  62
  63         def reset_after_request(self):
  64             Session.rollback()
  65             Session.remove()
  66
  67         @property
  68         def query(self):
  69             return Session.query
  70
  71 else:
  72     from sqlalchemy.orm import sessionmaker
  73
  74     class DatabaseManager(object):
  75         """
  76         Manage database connections.
  77
  78         The main method here is session_scope which can be used with a
  79         "with" statement to get a session that is properly torn down
  80         by the end of execution.
  81         """
  82         def __init__(self, engine):
  83             self.engine = engine
  84             self.Session = sessionmaker(bind=engine)
  85             set_models_as_attributes(self)
  86
  87         @contextmanager
  88         def session_scope(self):
  89             """
  90             This is a context manager, use like::
  91
  92               with dbmanager.session_scope() as request.db:
  93                   some_view(request)
  94             """
  95             session = self.Session()
  96
  97             #####################################
  98             # Functions to emulate DatabaseMaster
  99             #####################################
 100             def save(obj):
 101                 session.add(obj)
 102                 session.flush()
 103
 104             def check_session_clean():
 105                 # Is this implemented right?
 106                 for dummy in session:
 107                     _log.warn("STRANGE: There are elements in the sql session. "
 108                               "Please report this and help us track this down.")
 109                     break
 110
 111             def reset_after_request():
 112                 session.rollback()
 113                 session.remove()
 114
 115             # now attach
 116             session.save = save
 117             session.check_session_clean = check_session_clean
 118             session.reset_after_request = reset_after_request
 119
 120             set_models_as_attributes(session)
 121             #####################################
 122
 123             try:
 124                 yield session
 125             finally:
 126                 session.rollback()
 127                 session.close()
 128
 129
 130
 131 def load_models(app_config):
 132     import mediagoblin.db.models
 133
 134     for plugin in mg_globals.global_config.get('plugins', {}).keys():
 135         _log.debug("Loading %s.models", plugin)
 136         try:
 137             __import__(plugin + ".models")
 138         except ImportError as exc:
 139             _log.debug("Could not load {0}.models: {1}".format(
 140                 plugin,
 141                 exc))
 142
 143
 144 def _sqlite_fk_pragma_on_connect(dbapi_con, con_record):
 145     """Enable foreign key checking on each new sqlite connection"""
 146     dbapi_con.execute('pragma foreign_keys=on')
 147
 148
 149 def _sqlite_disable_fk_pragma_on_connect(dbapi_con, con_record):
 150     """
 151     Disable foreign key checking on each new sqlite connection
 152     (Good for migrations!)
 153     """
 154     dbapi_con.execute('pragma foreign_keys=off')
 155
 156
 157 def setup_connection_and_db_from_config(app_config, migrations=False, app=None):
 158     engine = create_engine(app_config['sql_engine'])
 159
 160     # @@: Maybe make a weak-ref so an engine can get garbage
 161     # collected?  Not that we expect to make a lot of MediaGoblinApp
 162     # instances in a single process...
 163     engine.app = app
 164
 165     # Enable foreign key checking for sqlite
 166     if app_config['sql_engine'].startswith('sqlite://'):
 167         if migrations:
 168             event.listen(engine, 'connect',
 169                          _sqlite_disable_fk_pragma_on_connect)
 170         else:
 171             event.listen(engine, 'connect', _sqlite_fk_pragma_on_connect)
 172
 173     # logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
 174
 175     if DISABLE_GLOBALS:
 176         return DatabaseManager(engine)
 177
 178     else:
 179         Session.configure(bind=engine)
 180
 181         return DatabaseMaster(engine)
 182
 183
 184 def check_db_migrations_current(db):
 185     pass