Source code for oriole_service.db

#
#                __   _,--="=--,_   __
#               /  \."    .-.    "./  \
#              /  ,/  _   : :   _  \/` \
#              \  `| /o\  :_:  /o\ |\__/
#               `-'| :="~` _ `~"=: |
#                  \`     (_)     `/
#           .-"-.   \      |      /   .-"-.
#    .-----{     }--|  /,.-'-.,\  |--{     }-----.
#     )    (_)_)_)  \_/`~-===-~`\_/  (_(_(_)    (
#    (                                           )
#     )                 Oriole-DB               (
#    (                  Eric.Zhou                )
#    '-------------------------------------------'
#

from weakref import WeakKeyDictionary

from nameko.extensions import DependencyProvider

from oriole.db import *


MyBase = declarative_base()


[docs]class Base(MyBase): __abstract__ = True def __init__(self, **kwargs): for attr in self.__mapper__.column_attrs: key = attr.key if key in kwargs: continue val = attr.columns[0].default if val and not callable(val.arg): kwargs[key] = val.arg super().__init__(**kwargs)
[docs]class Db(DependencyProvider): def __init__(self, Base, uri="database", read_only=False): self.base = Base self.uri = uri self.dbs = WeakKeyDictionary() self.read_only = read_only
[docs] def setup(self): self.bind = get_engine(self.container.config.get(self.uri)) if self.read_only: return self.base.metadata.create_all(self.bind)
[docs] def get_dependency(self, worker_ctx): session = get_session(self.bind) self.dbs[worker_ctx] = session return session
[docs] def worker_teardown(self, worker_ctx): session = self.dbs.pop(worker_ctx) session.close()
[docs]class Rs(DependencyProvider): def __init__(self, uri="datasets"): self.uri = uri
[docs] def setup(self): self.rs = get_redis(self.container.config.get(self.uri)) if self.rs: self.rs.current_ms_config = self.container.config
[docs] def get_dependency(self, worker_ctx): return self.rs