Merge pull request #61 from absent1706/fix/entity-zero-call

PR 51

.github/workflows/test-pr.yml

@@ -21,8 +21,20 @@ jobs:
         python -m pip install --upgrade pip setuptools wheel
         pip install -r requirements-dev.txt
         pip install "sqlalchemy<1.4"
-    - name: Run tests
+    - name: Run tests < sqlalchemy1.4
       run: |
+        nosetests --with-coverage --cover-inclusive --cover-package=sqlalchemy_mixins
+        export PYTHONPATH=.:$PYTHONPATH
+        python examples/activerecord.py
+        python examples/all_features.py
+        python examples/eagerload.py
+        python examples/repr.py
+        python examples/smartquery.py
+        python examples/serialize.py
+        python examples/timestamp.py
+    - name: Run tests >= sqlalchemy1.4
+      run: |
+        pip install -U sqlalchemy
         nosetests --with-coverage --cover-inclusive --cover-package=sqlalchemy_mixins
         export PYTHONPATH=.:$PYTHONPATH
         python examples/activerecord.py

sqlalchemy_mixins/smartquery.py

@@ -62,6 +62,17 @@ def _parse_path_and_make_aliases(entity, entity_path, attrs, aliases):
         aliases[path] = alias, relationship
         _parse_path_and_make_aliases(alias, path, nested_attrs, aliases)
 
+def _get_root_cls(query):
+    # sqlalchemy < 1.4.0
+    if hasattr(query, '_entity_zero'):
+        return query._entity_zero().class_
+
+    # sqlalchemy >= 1.4.0
+    else:
+        if hasattr(query, '_entity_from_pre_ent_zero'):
+            return query._entity_from_pre_ent_zero().class_
+    raise ValueError('Cannot get a root class from`{}`'
+                     .format(query))
 
 def smart_query(query, filters=None, sort_attrs=None, schema=None):
     """
@@ -83,8 +94,16 @@ def smart_query(query, filters=None, sort_attrs=None, schema=None):
     if not schema:
         schema = {}
 
+    # sqlalchemy >= 1.4.0, should probably a. check something else to determine if we need to convert
+    # AppenderQuery to a query, b. probably not hack it like this
     # noinspection PyProtectedMember
-    root_cls = query._entity_zero().class_  # for example, User or Post
+    if type(query).__name__ == 'AppenderQuery' and query._statement:
+        sess = query.session
+        # noinspection PyProtectedMember
+        query = query._statement
+        query.session = sess
+
+    root_cls = _get_root_cls(query)  # for example, User or Post
     attrs = list(filters.keys()) + \
         list(map(lambda s: s.lstrip(DESC_PREFIX), sort_attrs))
     aliases = OrderedDict({})

sqlalchemy_mixins/tests/test_smartquery.py

@@ -33,6 +33,13 @@ class User(BaseModel):
     # not to be a backref
     posts = sa.orm.relationship('Post')
     comments = sa.orm.relationship('Comment')
+    # below relationship will just return query (without executing)
+    # this query can be customized
+    # see http://docs.sqlalchemy.org/en/latest/orm/collections.html#dynamic-relationship
+    #
+    # we will use this relationship for demonstrating real-life example
+    #  of how smart_query() function works (see 3.2.2)
+    comments_ = sa.orm.relationship('Comment', lazy="dynamic")  # this will return query
 
 
 class Post(BaseModel):
@@ -217,7 +224,7 @@ def test_filterable_attributes(self):
 
         self.assertEqual(set(User.filterable_attributes),
                          {'id', 'name',  # normal columns
-                          'posts', 'comments'  # relations
+                          'posts', 'comments', 'comments_'  # relations
                           })
         self.assertNotIn('posts_viewonly', set(User.filterable_attributes))
 
@@ -785,6 +792,31 @@ def test_explicitly_set_in_schema_subqueryload(self):
         _ = res[0].post.comments
         self.assertEqual(self.query_count, 2)
 
+    def test_lazy_dynamic(self):
+        u1, u2, u3, p11, p12, p21, p22, cm11, cm12, cm21, cm22, cm_empty = \
+            self._seed()
+
+        schema = {
+            'post': {
+                'user': JOINED
+            }
+        }
+
+        user = sess.query(User).first()
+        # and we have initial query for his/her comments
+        #  (see User.comments_ relationship)
+        query = user.comments_
+        # now we just smartly apply all filters, sorts and eagerload. Perfect!
+        res = smart_query(query,
+                          filters={
+                              'post___public': True,
+                              'user__isnull': False
+                          },
+                          sort_attrs=['user___name', '-created_at'],
+                          schema=schema).all()
+
+        assert res[0] == cm21
+
     # TODO: implement below logic
     @nose.tools.raises(Exception)
     def test_override_eagerload_method_in_schema(self):