diff -r 000000000000 -r 17f6c240cd7b src/pyams_alchemy/engine.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/pyams_alchemy/engine.py Tue Mar 03 17:12:58 2015 +0100 @@ -0,0 +1,166 @@ +# +# Copyright (c) 2008-2015 Thierry Florac +# All Rights Reserved. +# +# This software is subject to the provisions of the Zope Public License, +# Version 2.1 (ZPL). A copy of the ZPL should accompany this distribution. +# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED +# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS +# FOR A PARTICULAR PURPOSE. +# + +__docformat__ = 'restructuredtext' + + +# import standard library +import sqlalchemy +from sqlalchemy.orm.scoping import scoped_session +from sqlalchemy.orm.session import sessionmaker + +# import interfaces +from pyams_alchemy.interfaces import REQUEST_SESSION_KEY, IAlchemyEngineUtility +from pyams_utils.interfaces.site import IOptionalUtility +from zope.lifecycleevent.interfaces import IObjectAddedEvent, IObjectRemovedEvent +from zope.schema.interfaces import IVocabularyFactory + +# import packages +from persistent import Persistent +from persistent.dict import PersistentDict +from pyams_utils.registry import query_utility +from pyams_utils.request import check_request, get_request_data, set_request_data +from pyramid.events import subscriber +from zope.componentvocabulary.vocabulary import UtilityVocabulary +from zope.container.contained import Contained +from zope.interface import implementer, provider +from zope.schema.fieldproperty import FieldProperty +from zope.schema.vocabulary import getVocabularyRegistry +from zope.sqlalchemy.datamanager import _SESSION_STATE, STATUS_READONLY, STATUS_ACTIVE, \ + ZopeTransactionExtension, join_transaction + + +@implementer(IAlchemyEngineUtility, IOptionalUtility) +class AlchemyEngineUtility(object): + """SQLAlchemy engine utility""" + + name = FieldProperty(IAlchemyEngineUtility['name']) + dsn = FieldProperty(IAlchemyEngineUtility['dsn']) + echo = FieldProperty(IAlchemyEngineUtility['echo']) + pool_size = FieldProperty(IAlchemyEngineUtility['pool_size']) + pool_recycle = FieldProperty(IAlchemyEngineUtility['pool_recycle']) + encoding = FieldProperty(IAlchemyEngineUtility['encoding']) + convert_unicode = FieldProperty(IAlchemyEngineUtility['convert_unicode']) + + def __init__(self, name='', dsn='', echo=False, pool_size=25, pool_recycle=-1, + encoding='utf-8', convert_unicode=False, **kwargs): + self.name = name + self.dsn = dsn + self.echo = echo + self.pool_size = pool_size + self.pool_recycle = pool_recycle + self.encoding = encoding + self.convert_unicode = convert_unicode + self.kw = PersistentDict() + self.kw.update(kwargs) + + def __setattr__(self, key, value): + super(AlchemyEngineUtility, self).__setattr__(key, value) + if (key != '_v_engine') and hasattr(self, '_v_engine'): + delattr(self, '_v_engine') + + def get_engine(self): + engine = getattr(self, '_v_engine', None) + if engine is not None: + return engine + kw = {} + kw.update(self.kw) + self._v_engine = sqlalchemy.create_engine(self.dsn, + echo=self.echo, + pool_size=self.pool_size, + pool_recycle=self.pool_recycle, + encoding=self.encoding, + convert_unicode=self.convert_unicode, + strategy='threadlocal', + **kw) + return self._v_engine + + +class PersistentAlchemyEngineUtility(Persistent, AlchemyEngineUtility, Contained): + """Persistent implementation of SQLAlchemy engine utility""" + + +@subscriber(IObjectAddedEvent, context_selector=IAlchemyEngineUtility) +def handle_added_engine(event): + """Register new SQLAlchemy engine when added""" + manager = event.newParent + manager.registerUtility(event.object, IAlchemyEngineUtility, name=event.object.name or '') + + +@subscriber(IObjectRemovedEvent, context_selector=IAlchemyEngineUtility) +def handle_removed_engine(event): + """Un-register an SQLAlchemy engine when deleted""" + manager = event.oldParent + manager.unregisterUtility(event.object, IAlchemyEngineUtility, name=event.object.name or '') + + +def get_engine(engine): + """Get engine matching given utility name""" + if isinstance(engine, str): + engine = query_utility(IAlchemyEngineUtility, name=engine) + if engine is not None: + return engine.get_engine() + + +def get_session(engine, join=True, status=STATUS_ACTIVE, request=None, alias=None, + twophase=True, use_zope_extension=True): + """Get a new SQLALchemy session + + Session is stored in request and in session storage + """ + if request is None: + request = check_request() + if not alias: + alias = engine + session_data = get_request_data(request, REQUEST_SESSION_KEY, {}) + session = session_data.get(alias) + if session is None: + _engine = get_engine(engine) + if use_zope_extension: + factory = scoped_session(sessionmaker(bind=_engine, + twophase=twophase, + extension=ZopeTransactionExtension())) + else: + factory = sessionmaker(bind=_engine, twophase=twophase) + session = factory() + if join: + join_transaction(session, initial_state=status) + if status != STATUS_READONLY: + _SESSION_STATE[id(session)] = session + if session is not None: + session_data[alias] = session + set_request_data(request, REQUEST_SESSION_KEY, session_data) + return session + + +def get_user_session(engine, join=True, status=STATUS_ACTIVE, request=None, alias=None, + twophase=True, use_zope_extension=True): + """Get a new SQLAlchemy session + + `engine` can be a session name or an already created session (in which case it's + returned as-is). + """ + if isinstance(engine, str): + session = get_session(engine, join, status, request, alias, twophase, use_zope_extension) + else: + session = engine + return session + + +@provider(IVocabularyFactory) +class EnginesVocabulary(UtilityVocabulary): + """SQLAlchemy engines vocabulary""" + + interface = IAlchemyEngineUtility + nameOnly = True + +getVocabularyRegistry().register('PyAMS SQLAlchemy engines', EnginesVocabulary)