142 lines
4.5 KiB
Python
142 lines
4.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
pyjeeves
|
|
~~~~~~~~~~~~~~~
|
|
|
|
Global objects
|
|
"""
|
|
from pyjeeves import logging, config
|
|
from weakref import WeakValueDictionary
|
|
from sqlalchemy import create_engine, orm
|
|
from sqlalchemy.orm import sessionmaker, scoped_session, Query, aliased
|
|
from sqlalchemy.orm.exc import UnmappedClassError
|
|
from sqlalchemy.ext.declarative import declarative_base, DeclarativeMeta
|
|
|
|
logger = logging.getLogger("PyJeeves." + __name__)
|
|
|
|
|
|
class BaseFilterQuery(Query):
|
|
def get(self, ident):
|
|
# Override get() so that the flag is always checked in the
|
|
# DB as opposed to pulling from the identity map. - this is optional.
|
|
return Query.get(self.populate_existing(), ident)
|
|
|
|
def __iter__(self):
|
|
return Query.__iter__(self.private())
|
|
|
|
def from_self(self, *ent):
|
|
# Override from_self() to automatically apply
|
|
# the criterion to. this works with count() and
|
|
# others.
|
|
return Query.from_self(self.private(), *ent)
|
|
|
|
def private(self):
|
|
# Fetch the model name and column list and apply model-specific base filters
|
|
mzero = self._mapper_zero()
|
|
|
|
if mzero:
|
|
# Sometimes a plain model class will be fetched instead of mzero
|
|
try:
|
|
model = mzero.class_
|
|
obj = mzero.class_
|
|
except Exception:
|
|
model = mzero.__class__
|
|
obj = mzero
|
|
|
|
if hasattr(model, '_base_filters'):
|
|
return self.enable_assertions(False).filter(model._base_filters(obj))
|
|
|
|
return self
|
|
|
|
|
|
class Model(object):
|
|
"""Baseclass for custom user models."""
|
|
|
|
#: the query class used. The :attr:`query` attribute is an instance
|
|
#: of this class. By default a :class:`BaseQuery` is used.
|
|
query_class = BaseFilterQuery
|
|
|
|
#: an instance of :attr:`query_class`. Can be used to query the
|
|
#: database for instances of this model.
|
|
query = None
|
|
|
|
|
|
class MetaBaseModel(DeclarativeMeta):
|
|
""" Define a metaclass for the BaseModel
|
|
Implement `__getitem__` for managing aliases """
|
|
|
|
def __init__(cls, *args):
|
|
super().__init__(*args)
|
|
cls.aliases = WeakValueDictionary()
|
|
|
|
def __getitem__(cls, key):
|
|
try:
|
|
alias = cls.aliases[key]
|
|
except KeyError:
|
|
alias = aliased(cls)
|
|
cls.aliases[key] = alias
|
|
return alias
|
|
|
|
|
|
class _QueryProperty(object):
|
|
|
|
def __init__(self, sa):
|
|
self.sa = sa
|
|
|
|
def __get__(self, obj, type):
|
|
try:
|
|
mapper = orm.class_mapper(type)
|
|
if mapper:
|
|
if type.__module__ == 'pyjeeves.models.raw':
|
|
return type.query_class(mapper, session=self.sa.raw_session())
|
|
else:
|
|
return type.query_class(mapper, session=self.sa.meta_session())
|
|
except UnmappedClassError:
|
|
return None
|
|
|
|
|
|
class DBConnector(object):
|
|
"""This class is used to control the SQLAlchemy integration"""
|
|
def __init__(self, enabled_sessions=['raw'], metadata=None):
|
|
logger.info("Creating engines and sessionmakers")
|
|
|
|
self.raw_session, self.meta_session = self.create_scoped_session(enabled_sessions)
|
|
self.Model = self.make_declarative_base(metadata)
|
|
# self.Query = Query
|
|
|
|
@property
|
|
def metadata(self):
|
|
"""Returns the metadata"""
|
|
return self.Model.metadata
|
|
|
|
# @property
|
|
# def _config(self):
|
|
# """Returns the configuration"""
|
|
# return config()
|
|
|
|
def make_declarative_base(self, metadata=None):
|
|
"""Creates the declarative base."""
|
|
base = declarative_base(cls=Model, name='Model',
|
|
metadata=metadata,
|
|
metaclass=MetaBaseModel)
|
|
base.query = _QueryProperty(self)
|
|
return base
|
|
|
|
def create_scoped_session(self, sessions=[]):
|
|
RawSession, MetaSession = None, None
|
|
if 'raw' in sessions:
|
|
raw_engine = create_engine(
|
|
'mssql+pymssql://{user}:{pw}@{host}:{port}/{db}?charset=utf8'.format(
|
|
**config.config['databases']['raw']),
|
|
implicit_returning=False)
|
|
|
|
RawSession = scoped_session(sessionmaker(bind=raw_engine))
|
|
|
|
if 'meta' in sessions:
|
|
meta_engine = create_engine(
|
|
'mysql+pymysql://{user}:{pw}@{host}:{port}/{db}?charset=utf8mb4'.format(
|
|
**config.config['databases']['meta']))
|
|
|
|
MetaSession = scoped_session(sessionmaker(bind=meta_engine))
|
|
|
|
return RawSession, MetaSession
|