Added sort support to SQLAlchemyConnectionField

Author: CaselIT

examples/flask_sqlalchemy/schema.py

@@ -1,6 +1,6 @@
 import graphene
 from graphene import relay
-from graphene_sqlalchemy import SQLAlchemyConnectionField, SQLAlchemyObjectType
+from graphene_sqlalchemy import SQLAlchemyConnectionField, SQLAlchemyObjectType, utils
 from models import Department as DepartmentModel
 from models import Employee as EmployeeModel
 from models import Role as RoleModel
@@ -27,11 +27,17 @@ class Meta:
         interfaces = (relay.Node, )
 
 
+SortEnumEmployee = utils.sort_enum_for_model(
+    EmployeeModel, 'SortEnumEmployee',
+    lambda c, d: c.upper() + ('_ASC' if d else '_DESC'))
+
+
 class Query(graphene.ObjectType):
     node = relay.Node.Field()
-    all_employees = SQLAlchemyConnectionField(Employee)
-    all_roles = SQLAlchemyConnectionField(Role)
-    role = graphene.Field(Role)
+    all_employees = SQLAlchemyConnectionField(
+        Employee, sort=graphene.Argument(SortEnumEmployee, default_value=EmployeeModel.id))
+    all_roles = SQLAlchemyConnectionField(Role, sort=utils.sort_argument_for_model(RoleModel))
+    all_departments = SQLAlchemyConnectionField(Department)
 
 
 schema = graphene.Schema(query=Query, types=[Department, Employee, Role])

graphene_sqlalchemy/fields.py

@@ -1,3 +1,4 @@
+from collections import Iterable
 from functools import partial
 
 from sqlalchemy.orm.query import Query
@@ -16,8 +17,11 @@ def model(self):
         return self.type._meta.node._meta.model
 
     @classmethod
-    def get_query(cls, model, info, **args):
-        return get_query(model, info.context)
+    def get_query(cls, model, info, sort=None, **args):
+        query = get_query(model, info.context)
+        if sort is not None:
+            query = query.order_by(*sort) if isinstance(sort, Iterable) else query.order_by(sort)
+        return query
 
     @property
     def type(self):

graphene_sqlalchemy/utils.py

@@ -1,4 +1,6 @@
+from graphene import Argument, Enum, List
 from sqlalchemy.exc import ArgumentError
+from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import class_mapper, object_mapper
 from sqlalchemy.orm.exc import UnmappedClassError, UnmappedInstanceError
 
@@ -34,3 +36,51 @@ def is_mapped_instance(cls):
         return False
     else:
         return True
+
+
+def _symbol_name(column_name, is_asc):
+    return column_name + ('_asc' if is_asc else '_desc')
+
+
+def _sort_enum_for_model(cls, name=None, symbol_name=_symbol_name):
+    name = name or cls.__name__ + 'SortEnum'
+    items = []
+    default = []
+    for column in inspect(cls).columns.values():
+        asc = symbol_name(column.name, True), column.asc()
+        desc = symbol_name(column.name, False), column.desc()
+        if column.primary_key:
+            default.append(asc[1])
+        items.extend((asc, desc))
+    return Enum(name, items), default
+
+
+def sort_enum_for_model(cls, name=None, symbol_name=_symbol_name):
+    '''Create Graphene Enum for sorting a SQLAlchemy class query
+
+    Parameters
+    - cls : Sqlalchemy model class
+        Model used to create the sort enumerator
+    - name : str, optional, default None
+        Name to use for the enumerator. If not provided it will be set to `cls.__name__ + 'SortEnum'`
+    - symbol_name : function, optional, default `_symbol_name`
+        Function which takes the column name and a boolean indicating if the sort direction is ascending,
+        and returns the symbol name for the current column and sort direction.
+        The default function will create, for a column named 'foo', the symbols 'foo_asc' and 'foo_desc'
+
+    Returns
+    - Enum
+        The Graphene enumerator
+    '''
+    enum, _ = _sort_enum_for_model(cls, name, symbol_name)
+    return enum
+
+
+def sort_argument_for_model(cls, has_default=True):
+    '''Returns an Graphene argument for the sort field that accepts a list of sorting directions for a model.
+    If `has_default` is True (the default) it will sort the result by the primary key(s)
+    '''
+    enum, default = _sort_enum_for_model(cls)
+    if not has_default:
+        default = None
+    return Argument(List(enum), default_value=default)