Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement custom query and model class options #328

Merged
merged 3 commits into from
Oct 23, 2015
Merged

Implement custom query and model class options #328

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
155656a
Implement custom query and model class options
justanr May 8, 2015
f1de1e2
Add custom query class to session options.
justanr Aug 10, 2015
3ff4c64
Respect existing query class on models
justanr Aug 11, 2015
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 19 additions & 15 deletions flask_sqlalchemy/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,34 +68,34 @@ def _make_table(*args, **kwargs):
return _make_table


def _set_default_query_class(d):
def _set_default_query_class(d, cls):
if 'query_class' not in d:
d['query_class'] = BaseQuery
d['query_class'] = cls


def _wrap_with_default_query_class(fn):
def _wrap_with_default_query_class(fn, cls):
@functools.wraps(fn)
def newfn(*args, **kwargs):
_set_default_query_class(kwargs)
_set_default_query_class(kwargs, cls)
if "backref" in kwargs:
backref = kwargs['backref']
if isinstance(backref, string_types):
backref = (backref, {})
_set_default_query_class(backref[1])
_set_default_query_class(backref[1], cls)
return fn(*args, **kwargs)
return newfn


def _include_sqlalchemy(obj):
def _include_sqlalchemy(obj, cls):
for module in sqlalchemy, sqlalchemy.orm:
for key in module.__all__:
if not hasattr(obj, key):
setattr(obj, key, getattr(module, key))
# Note: obj.Table does not attempt to be a SQLAlchemy Table class.
obj.Table = _make_table(obj)
obj.relationship = _wrap_with_default_query_class(obj.relationship)
obj.relation = _wrap_with_default_query_class(obj.relation)
obj.dynamic_loader = _wrap_with_default_query_class(obj.dynamic_loader)
obj.relationship = _wrap_with_default_query_class(obj.relationship, cls)
obj.relation = _wrap_with_default_query_class(obj.relation, cls)
obj.dynamic_loader = _wrap_with_default_query_class(obj.dynamic_loader, cls)
obj.event = event


Expand Down Expand Up @@ -730,19 +730,23 @@ class User(db.Model):
naming conventions among other, non-trivial things.
"""

def __init__(self, app=None, use_native_unicode=True, session_options=None, metadata=None):
def __init__(self, app=None, use_native_unicode=True, session_options=None,
metadata=None, query_class=BaseQuery, model_class=Model):

if session_options is None:
session_options = {}

session_options.setdefault('scopefunc', connection_stack.__ident_func__)
session_options.setdefault('query_cls', query_class)
self.use_native_unicode = use_native_unicode
self.session = self.create_scoped_session(session_options)
self.Model = self.make_declarative_base(metadata)
self.Query = BaseQuery
self.Model = self.make_declarative_base(model_class, metadata)
# ugly hack, is there a better way?
self.Model.query_class = query_class
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not happy with this line. I considered moving it to make_declarative_base but I'm not 100%. Should we respect the query class on the model, leave this as is/move to make_declarative_base, add a toggle for replacing the query class or something else?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Everything dealing with the Model should be in one place. Just do something like if query_class is not None and getattr(Model, 'query_class', None) is None: Model.query_class = query_class to avoid overriding it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a check to respect the existing query_class attribute on models, and this potential replacement was moved into SQLAlchemy.make_declarative_base. However, I'm wondering if this'll cause confusion with some people. Providing a model base with the desired query_class will work though.

self.Query = query_class
self._engine_lock = Lock()
self.app = app
_include_sqlalchemy(self)
_include_sqlalchemy(self, query_class)

if app is not None:
self.init_app(app)
Expand Down Expand Up @@ -770,9 +774,9 @@ def create_session(self, options):
"""
return SignallingSession(self, **options)

def make_declarative_base(self, metadata=None):
def make_declarative_base(self, model_class, metadata=None):
"""Creates the declarative base."""
base = declarative_base(cls=Model, name='Model',
base = declarative_base(cls=model_class, name='Model',
metadata=metadata,
metaclass=_BoundDeclarativeMeta)
base.query = _QueryProperty(self)
Expand Down
54 changes: 52 additions & 2 deletions test_sqlalchemy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,18 +467,68 @@ def test_default_query_class(self):

class Parent(db.Model):
id = db.Column(db.Integer, primary_key=True)
children = db.relationship("Child", backref = "parents", lazy='dynamic')
children = db.relationship("Child", backref = "parent", lazy='dynamic')

class Child(db.Model):
id = db.Column(db.Integer, primary_key=True)
parent_id = db.Column(db.Integer, db.ForeignKey('parent.id'))

p = Parent()
c = Child()
c.parent = p

self.assertEqual(type(Parent.query), sqlalchemy.BaseQuery)
self.assertEqual(type(Child.query), sqlalchemy.BaseQuery)
self.assertTrue(isinstance(p.children, sqlalchemy.BaseQuery))
#self.assertTrue(isinstance(c.parents, sqlalchemy.BaseQuery))
self.assertTrue(isinstance(db.session.query(Parent), sqlalchemy.BaseQuery))


class CustomQueryClassTestCase(unittest.TestCase):

def test_custom_query_class(self):
class CustomQueryClass(sqlalchemy.BaseQuery):
pass

app = flask.Flask(__name__)
app.config['SQLALCHEMY_ENGINE'] = 'sqlite://'
app.config['TESTING'] = True
db = sqlalchemy.SQLAlchemy(app, query_class=CustomQueryClass)

class Parent(db.Model):
id = db.Column(db.Integer, primary_key=True)
children = db.relationship("Child", backref = "parent", lazy='dynamic')

class Child(db.Model):
id = db.Column(db.Integer, primary_key=True)
parent_id = db.Column(db.Integer, db.ForeignKey('parent.id'))

p = Parent()
c = Child()
c.parent = p

self.assertEqual(type(Parent.query), CustomQueryClass)
self.assertEqual(type(Child.query), CustomQueryClass)
self.assertTrue(isinstance(p.children, CustomQueryClass))
self.assertEqual(db.Query, CustomQueryClass)
self.assertEqual(db.Model.query_class, CustomQueryClass)
self.assertTrue(isinstance(db.session.query(Parent), CustomQueryClass))


class CustomModelClassTestCase(unittest.TestCase):

def test_custom_query_class(self):
class CustomModelClass(sqlalchemy.Model):
pass

app = flask.Flask(__name__)
app.config['SQLALCHEMY_ENGINE'] = 'sqlite://'
app.config['TESTING'] = True
db = sqlalchemy.SQLAlchemy(app, model_class=CustomModelClass)

class SomeModel(db.Model):
id = db.Column(db.Integer, primary_key=True)

self.assertTrue(isinstance(SomeModel(), CustomModelClass))

class SQLAlchemyIncludesTestCase(unittest.TestCase):

Expand Down