From 2f8359c6f25768805431c80c74e5ec4213c2b2a6 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 30 Mar 2018 08:05:37 -0700
Subject: [PATCH] Work in progress with failing tests

---
 datasette/app.py   |  5 +++++
 datasette/utils.py | 12 ++++++++++++
 tests/fixtures.py  | 37 ++++++++++++++++++++++++++++++++++++-
 tests/test_api.py  | 37 ++++++++++++++++++++++++++++++++++++-
 4 files changed, 89 insertions(+), 2 deletions(-)

diff --git a/datasette/app.py b/datasette/app.py
index ef3fde93..38f16e49 100644
--- a/datasette/app.py
+++ b/datasette/app.py
@@ -602,6 +602,11 @@ class TableView(RowTableShared):
         if where_clauses:
             where_clause = 'where {} '.format(' and '.join(where_clauses))
 
+        # Allow for custom sort order
+        sort = special_args.get('_sort')
+        if sort:
+            order_by = sort
+
         if order_by:
             order_by = 'order by {} '.format(order_by)
 
diff --git a/datasette/utils.py b/datasette/utils.py
index fb132f08..9860355c 100644
--- a/datasette/utils.py
+++ b/datasette/utils.py
@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from collections import namedtuple
 import base64
 import hashlib
 import json
@@ -12,6 +13,17 @@ import shutil
 import urllib
 
 
+Sort = namedtuple('Sort', ['column', 'desc', 'nulls_last'])
+sort_mapping = {
+    '_sort': Sort('_sort', False, False),
+    '_sort_asc': Sort('_sort', False, False),
+    '_sort_desc': Sort('_sort_desc', True, False),
+    '_sort_nulls_last': Sort('_sort_nulls_last', False, True),
+    '_sort_asc_nulls_last': Sort('_sort_nulls_last', False, True),
+    '_sort_desc_nulls_last': Sort('_sort_desc_nulls_last', True, True),
+}
+
+
 def compound_pks_from_path(path):
     return [
         urllib.parse.unquote_plus(b) for b in path.split(',')
diff --git a/tests/fixtures.py b/tests/fixtures.py
index 3d22c289..892b24e1 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -1,6 +1,7 @@
 from datasette.app import Datasette
 import itertools
 import os
+import random
 import sqlite3
 import string
 import tempfile
@@ -33,6 +34,25 @@ def generate_compound_rows(num):
         yield a, b, c, '{}-{}-{}'.format(a, b, c)
 
 
+def generate_sortable_rows(num):
+    rand = random.Random(42)
+    for a, b in itertools.islice(
+        itertools.product(string.ascii_lowercase, repeat=2), num
+    ):
+        yield {
+            'pk1': a,
+            'pk2': b,
+            'content': '{}-{}'.format(a, b),
+            'sortable': rand.randint(-100, 100),
+            'sortable_with_nulls': rand.choice([
+                None, rand.random(), rand.random()
+            ]),
+            'sortable_with_nulls_2': rand.choice([
+                None, rand.random(), rand.random()
+            ]),
+        }
+
+
 METADATA = {
     'title': 'Datasette Title',
     'description': 'Datasette Description',
@@ -69,7 +89,6 @@ CREATE TABLE compound_primary_key (
 
 INSERT INTO compound_primary_key VALUES ('a', 'b', 'c');
 
-
 CREATE TABLE compound_three_primary_keys (
   pk1 varchar(30),
   pk2 varchar(30),
@@ -78,6 +97,15 @@ CREATE TABLE compound_three_primary_keys (
   PRIMARY KEY (pk1, pk2, pk3)
 );
 
+CREATE TABLE sortable (
+  pk1 varchar(30),
+  pk2 varchar(30),
+  content text,
+  sortable integer,
+  sortable_with_nulls real,
+  sortable_with_nulls_2 real,
+  PRIMARY KEY (pk1, pk2)
+);
 
 CREATE TABLE no_primary_key (
   content text,
@@ -134,4 +162,11 @@ CREATE VIEW simple_view AS
     'INSERT INTO compound_three_primary_keys VALUES ("{a}", "{b}", "{c}", "{content}");'.format(
         a=a, b=b, c=c, content=content
     ) for a, b, c, content in generate_compound_rows(1001)
+]) + '\n'.join([
+    '''INSERT INTO sortable VALUES (
+        "{pk1}", "{pk2}", "{content}", {sortable},
+        {sortable_with_nulls}, {sortable_with_nulls_2});
+    '''.format(
+        **row
+    ).replace('None', 'null') for row in generate_sortable_rows(201)
 ])
diff --git a/tests/test_api.py b/tests/test_api.py
index cff7dadc..476bf05e 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1,6 +1,7 @@
 from .fixtures import (
     app_client,
     generate_compound_rows,
+    generate_sortable_rows,
 )
 import pytest
 
@@ -13,7 +14,7 @@ def test_homepage(app_client):
     assert response.json.keys() == {'test_tables': 0}.keys()
     d = response.json['test_tables']
     assert d['name'] == 'test_tables'
-    assert d['tables_count'] == 8
+    assert d['tables_count'] == 9
 
 
 def test_database_page(app_client):
@@ -99,6 +100,16 @@ def test_database_page(app_client):
             'outgoing': [],
         },
         'label_column': None,
+    }, {
+        'columns': [
+            'pk1', 'pk2', 'content', 'sortable', 'sortable_with_nulls',
+            'sortable_with_nulls_2'
+        ],
+        'name': 'sortable',
+        'count': 201,
+        'hidden': False,
+        'foreign_keys': {'incoming': [], 'outgoing': []},
+        'label_column': None,
     }, {
         'columns': ['pk', 'content'],
         'name': 'table/with/slashes.csv',
@@ -247,6 +258,30 @@ def test_paginate_compound_keys_with_extra_filters(app_client):
     assert expected == [f['content'] for f in fetched]
 
 
+@pytest.mark.parametrize('query_string,sort_key', [
+    ('_sort=sortable', lambda row: row['sortable']),
+    ('_sort_desc=sortable', lambda row: -row['sortable']),
+])
+def test_sortable(app_client, query_string, sort_key):
+    path = '/test_tables/sortable.jsono?{}'.format(query_string)
+    fetched = []
+    page = 0
+    while path:
+        page += 1
+        assert page < 100
+        response = app_client.get(path, gather_request=False)
+        fetched.extend(response.json['rows'])
+        path = response.json['next_url']
+    assert 5 == page
+    expected = list(generate_sortable_rows(201))
+    expected.sort(key=sort_key)
+    assert [
+        r['content'] for r in expected
+    ] == [
+        r['content'] for r in fetched
+    ]
+
+
 @pytest.mark.parametrize('path,expected_rows', [
     ('/test_tables/simple_primary_key.json?content=hello', [
         ['1', 'hello'],