Created
February 7, 2016 04:37
-
-
Save kennethreitz/9ccb1bf8434a51aa43be to your computer and use it in GitHub Desktop.
Some code for using schemas as your declarative base with SQL-Alchemy, along with support for Celery's forking.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from sqlalchemy import create_engine, MetaData, Table | |
from sqlalchemy.orm import sessionmaker | |
from sqlalchemy.ext.declarative import declarative_base | |
from multiprocessing.util import register_after_fork | |
Base = declarative_base() | |
_ENGINES = {} | |
_SESSIONS = {} | |
DATABASE_URL='insert ur here' | |
# https://github.com/celery/celery/blob/master/celery/backends/database/session.py#L27-L46 | |
class _after_fork(object): | |
registered = False | |
def __call__(self): | |
self.registered = False # child must reregister | |
for engine in list(_ENGINES.values()): | |
engine.dispose() | |
_ENGINES.clear() | |
_SESSIONS.clear() | |
after_fork = _after_fork() | |
def get_engine(dburi, **kwargs): | |
try: | |
return _ENGINES[dburi] | |
except KeyError: | |
engine = _ENGINES[dburi] = create_engine(dburi, **kwargs) | |
after_fork.registered = True | |
register_after_fork(after_fork, after_fork) | |
return engine | |
def create_session(dburi, short_lived_sessions=False, **kwargs): | |
engine = get_engine(dburi, **kwargs) | |
if short_lived_sessions or dburi not in _SESSIONS: | |
_SESSIONS[dburi] = sessionmaker(bind=engine) | |
return engine, _SESSIONS[dburi] | |
engine = create_engine(DATABASE_URL) | |
metadata = MetaData(bind=engine) | |
session_maker = create_session(DATABASE_URL)[1] | |
class SessionPropertyMixin(object): | |
def save(self, session=None, commit=True, close=False): | |
if session: | |
session.add(self) | |
session.commit() | |
if close: | |
session.close() | |
# Reflect each database table we need to use, via schemas. | |
class Slugs(Base, SessionPropertyMixin): | |
__table__ = Table('slugs', metadata, autoload=True) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment