Use custom query class

A custom query class allows to add more methods on
queries (read: "cursors").

This custom query class especially adds a .sort with a
calling convention exactly like the mongo one.
Makes a lot of existing code happy!

1 files changed, 21 insertions, 1 deletions

diff --git a/mediagoblin/db/sql/base.py b/mediagoblin/db/sql/base.py
index 40140327..082d498e 100644
--- a/mediagoblin/db/sql/base.py
+++ b/mediagoblin/db/sql/base.py
@@ -1,7 +1,27 @@
 from sqlalchemy.orm import scoped_session, sessionmaker, object_session
+from sqlalchemy.orm.query import Query
+from sqlalchemy.sql.expression import desc
+from mediagoblin.db.sql.fake import DESCENDING
 
 
-Session = scoped_session(sessionmaker())
+def _get_query_model(query):
+    cols = query.column_descriptions
+    assert len(cols) == 1, "These functions work only on simple queries"
+    return cols[0]["type"]
+
+
+class GMGQuery(Query):
+    def sort(self, key, direction):
+        key_col = getattr(_get_query_model(self), key)
+        if direction is DESCENDING:
+            key_col = desc(key_col)
+        return self.order_by(key_col)
+
+    def skip(self, amount):
+        return self.offset(amount)
+
+
+Session = scoped_session(sessionmaker(query_cls=GMGQuery))
 
 
 def _fix_query_dict(query_dict):