1 #

     2 # Copyright (c) 2008-2015 Thierry Florac <tflorac AT ulthar.net>

     3 # All Rights Reserved.

     4 #

     5 # This software is subject to the provisions of the Zope Public License,

     6 # Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.

     7 # THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED

     8 # WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED

     9 # WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS

    10 # FOR A PARTICULAR PURPOSE.

    11 #

    12 

    13 __docformat__ = 'restructuredtext'

    14 

    15 

    16 # import standard library

    17 import sqlalchemy

    18 from sqlalchemy.orm.scoping import scoped_session

    19 from sqlalchemy.orm.session import sessionmaker

    20 

    21 # import interfaces

    22 from pyams_alchemy.interfaces import REQUEST_SESSION_KEY, IAlchemyEngineUtility

    23 from pyams_utils.interfaces.site import IOptionalUtility

    24 from zope.lifecycleevent.interfaces import IObjectAddedEvent, IObjectRemovedEvent

    25 from zope.schema.interfaces import IVocabularyFactory

    26 

    27 # import packages

    28 from persistent import Persistent

    29 from persistent.dict import PersistentDict

    30 from pyams_utils.registry import query_utility

    31 from pyams_utils.request import check_request, get_request_data, set_request_data

    32 from pyramid.events import subscriber

    33 from zope.componentvocabulary.vocabulary import UtilityVocabulary

    34 from zope.container.contained import Contained

    35 from zope.interface import implementer, provider

    36 from zope.schema.fieldproperty import FieldProperty

    37 from zope.schema.vocabulary import getVocabularyRegistry

    38 from zope.sqlalchemy.datamanager import _SESSION_STATE, STATUS_READONLY, STATUS_ACTIVE, \

    39     ZopeTransactionExtension, join_transaction

    40 

    41 

    42 @implementer(IAlchemyEngineUtility, IOptionalUtility)

    43 class AlchemyEngineUtility(object):

    44     """SQLAlchemy engine utility"""

    45 

    46     name = FieldProperty(IAlchemyEngineUtility['name'])

    47     dsn = FieldProperty(IAlchemyEngineUtility['dsn'])

    48     echo = FieldProperty(IAlchemyEngineUtility['echo'])

    49     pool_size = FieldProperty(IAlchemyEngineUtility['pool_size'])

    50     pool_recycle = FieldProperty(IAlchemyEngineUtility['pool_recycle'])

    51     encoding = FieldProperty(IAlchemyEngineUtility['encoding'])

    52     convert_unicode = FieldProperty(IAlchemyEngineUtility['convert_unicode'])

    53 

    54     def __init__(self, name='', dsn='', echo=False, pool_size=25, pool_recycle=-1,

    55                  encoding='utf-8', convert_unicode=False, **kwargs):

    56         self.name = name

    57         self.dsn = dsn

    58         self.echo = echo

    59         self.pool_size = pool_size

    60         self.pool_recycle = pool_recycle

    61         self.encoding = encoding

    62         self.convert_unicode = convert_unicode

    63         self.kw = PersistentDict()

    64         self.kw.update(kwargs)

    65 

    66     def __setattr__(self, key, value):

    67         super(AlchemyEngineUtility, self).__setattr__(key, value)

    68         if (key != '_v_engine') and hasattr(self, '_v_engine'):

    69             delattr(self, '_v_engine')

    70 

    71     def get_engine(self):

    72         engine = getattr(self, '_v_engine', None)

    73         if engine is not None:

    74             return engine

    75         kw = {}

    76         kw.update(self.kw)

    77         self._v_engine = sqlalchemy.create_engine(self.dsn,

    78                                                   echo=self.echo,

    79                                                   pool_size=self.pool_size,

    80                                                   pool_recycle=self.pool_recycle,

    81                                                   encoding=self.encoding,

    82                                                   convert_unicode=self.convert_unicode,

    83                                                   strategy='threadlocal',

    84                                                   **kw)

    85         return self._v_engine

    86 

    87 

    88 class PersistentAlchemyEngineUtility(Persistent, AlchemyEngineUtility, Contained):

    89     """Persistent implementation of SQLAlchemy engine utility"""

    90 

    91 

    92 @subscriber(IObjectAddedEvent, context_selector=IAlchemyEngineUtility)

    93 def handle_added_engine(event):

    94     """Register new SQLAlchemy engine when added"""

    95     manager = event.newParent

    96     manager.registerUtility(event.object, IAlchemyEngineUtility, name=event.object.name or '')

    97 

    98 

    99 @subscriber(IObjectRemovedEvent, context_selector=IAlchemyEngineUtility)

   100 def handle_removed_engine(event):

   101     """Un-register an SQLAlchemy engine when deleted"""

   102     manager = event.oldParent

   103     manager.unregisterUtility(event.object, IAlchemyEngineUtility, name=event.object.name or '')

   104 

   105 

   106 def get_engine(engine):

   107     """Get engine matching given utility name"""

   108     if isinstance(engine, str):

   109         engine = query_utility(IAlchemyEngineUtility, name=engine)

   110         if engine is not None:

   111             return engine.get_engine()

   112 

   113 

   114 def get_session(engine, join=True, status=STATUS_ACTIVE, request=None, alias=None,

   115                 twophase=True, use_zope_extension=True):

   116     """Get a new SQLALchemy session

   117 

   118     Session is stored in request and in session storage

   119     """

   120     if request is None:

   121         request = check_request()

   122     if not alias:

   123         alias = engine

   124     session_data = get_request_data(request, REQUEST_SESSION_KEY, {})

   125     session = session_data.get(alias)

   126     if session is None:

   127         _engine = get_engine(engine)

   128         if use_zope_extension:

   129             factory = scoped_session(sessionmaker(bind=_engine,

   130                                                   twophase=twophase,

   131                                                   extension=ZopeTransactionExtension()))

   132         else:

   133             factory = sessionmaker(bind=_engine, twophase=twophase)

   134         session = factory()

   135         if join:

   136             join_transaction(session, initial_state=status)

   137         if status != STATUS_READONLY:

   138             _SESSION_STATE[id(session)] = session

   139         if session is not None:

   140             session_data[alias] = session

   141             set_request_data(request, REQUEST_SESSION_KEY, session_data)

   142     return session

   143 

   144 

   145 def get_user_session(engine, join=True, status=STATUS_ACTIVE, request=None, alias=None,

   146                      twophase=True, use_zope_extension=True):

   147     """Get a new SQLAlchemy session

   148 

   149     `engine` can be a session name or an already created session (in which case it's

   150     returned as-is).

   151     """

   152     if isinstance(engine, str):

   153         session = get_session(engine, join, status, request, alias, twophase, use_zope_extension)

   154     else:

   155         session = engine

   156     return session

   157 

   158 

   159 @provider(IVocabularyFactory)

   160 class EnginesVocabulary(UtilityVocabulary):

   161     """SQLAlchemy engines vocabulary"""

   162 

   163     interface = IAlchemyEngineUtility

   164     nameOnly = True

   165 

   166 getVocabularyRegistry().register('PyAMS SQLAlchemy engines', EnginesVocabulary)