From f643f7aee1b3d29df82b93a6560887a35b7a2f13 Mon Sep 17 00:00:00 2001
From: Simon Willison <simonw@eventbrite.com>
Date: Tue, 24 Oct 2017 07:58:41 -0700
Subject: [PATCH] base64 encode bytestrings from DB in JSON

Fixes #29
---
 app.py          | 29 +++++++++++++++++++++++++++--
 test_helpers.py | 19 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/app.py b/app.py
index afb5216d..b5a1445d 100644
--- a/app.py
+++ b/app.py
@@ -8,6 +8,7 @@ from pathlib import Path
 from functools import wraps
 import urllib.parse
 import json
+import base64
 import hashlib
 import sys
 
@@ -99,8 +100,15 @@ class BaseView(HTTPMethodView):
                 'error': str(e),
             }
         if as_json:
-            r = response.json(data)
-            r.headers['Access-Control-Allow-Origin'] = '*'
+            r = response.HTTPResponse(
+                json.dumps(
+                    data, cls=CustomJSONEncoder
+                ),
+                content_type='application/json',
+                headers={
+                    'Access-Control-Allow-Origin': '*'
+                }
+            )
         else:
             context = {**data, **dict(
                 extra_template_data()
@@ -159,6 +167,7 @@ class TableView(BaseView):
         rows = conn.execute('select * from {} limit 20'.format(table))
         columns = [r[0] for r in rows.description]
         pks = pks_for_table(conn, table)
+        rows = list(rows)
         return {
             'database': name,
             'table': table,
@@ -266,6 +275,22 @@ def path_from_row_pks(row, pks):
     return ','.join(bits)
 
 
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, sqlite3.Row):
+            return dict(obj)
+        if isinstance(obj, bytes):
+            # Does it encode to utf8?
+            try:
+                return obj.decode('utf8')
+            except UnicodeDecodeError:
+                return {
+                    '$base64': True,
+                    'encoded': base64.b64encode(obj).decode('latin1'),
+                }
+        return json.JSONEncoder.default(self, obj)
+
+
 if __name__ == '__main__':
     if '--build' in sys.argv:
         ensure_build_metadata(True)
diff --git a/test_helpers.py b/test_helpers.py
index 2149dadc..42e90166 100644
--- a/test_helpers.py
+++ b/test_helpers.py
@@ -1,6 +1,7 @@
 import app
 import pytest
 import sqlite3
+import json
 
 
 @pytest.mark.parametrize('path,expected', [
@@ -45,3 +46,21 @@ def test_pks_for_table(sql, table, expected_keys):
 def test_path_from_row_pks(row, pks, expected_path):
     actual_path = app.path_from_row_pks(row, pks)
     assert expected_path == actual_path
+
+
+@pytest.mark.parametrize('obj,expected', [
+    ({
+        'Description': 'Soft drinks',
+        'Picture': b"\x15\x1c\x02\xc7\xad\x05\xfe",
+        'CategoryID': 1,
+    }, """
+        {"CategoryID": 1, "Description": "Soft drinks", "Picture": {"$base64": true, "encoded": "FRwCx60F/g=="}}
+    """.strip()),
+])
+def test_custom_json_encoder(obj, expected):
+    actual = json.dumps(
+        obj,
+        cls=app.CustomJSONEncoder,
+        sort_keys=True
+    )
+    assert expected == actual