From 872dae1e1a1511e2edfb9d7ddf6ea5096c11d5c3 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 22 Dec 2023 15:08:11 -0800
Subject: [PATCH] Fix for CSV labels=on missing foreign key bug, closes #2214

---
 datasette/views/base.py | 14 ++++++++------
 tests/test_csv.py       | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/datasette/views/base.py b/datasette/views/base.py
index db08557e..e59fd683 100644
--- a/datasette/views/base.py
+++ b/datasette/views/base.py
@@ -553,16 +553,18 @@ async def stream_csv(datasette, fetch_data, request, database):
                                 if cell is None:
                                     new_row.extend(("", ""))
                                 else:
-                                    assert isinstance(cell, dict)
-                                    new_row.append(cell["value"])
-                                    new_row.append(cell["label"])
+                                    if not isinstance(cell, dict):
+                                        new_row.extend((cell, ""))
+                                    else:
+                                        new_row.append(cell["value"])
+                                        new_row.append(cell["label"])
                             else:
                                 new_row.append(cell)
                         await writer.writerow(new_row)
-            except Exception as e:
-                sys.stderr.write("Caught this error: {}\n".format(e))
+            except Exception as ex:
+                sys.stderr.write("Caught this error: {}\n".format(ex))
                 sys.stderr.flush()
-                await r.write(str(e))
+                await r.write(str(ex))
                 return
         await limited_writer.write(postamble)
 
diff --git a/tests/test_csv.py b/tests/test_csv.py
index ed83d685..9f772f89 100644
--- a/tests/test_csv.py
+++ b/tests/test_csv.py
@@ -1,3 +1,4 @@
+from datasette.app import Datasette
 from bs4 import BeautifulSoup as Soup
 import pytest
 from .fixtures import (  # noqa
@@ -95,6 +96,40 @@ async def test_table_csv_with_nullable_labels(ds_client):
     assert response.text == EXPECTED_TABLE_WITH_NULLABLE_LABELS_CSV
 
 
+@pytest.mark.asyncio
+async def test_table_csv_with_invalid_labels():
+    # https://github.com/simonw/datasette/issues/2214
+    ds = Datasette()
+    await ds.invoke_startup()
+    db = ds.add_memory_database("db_2214")
+    await db.execute_write_script(
+        """
+        create table t1 (id integer primary key, name text);
+        insert into t1 (id, name) values (1, 'one');
+        insert into t1 (id, name) values (2, 'two');
+        create table t2 (textid text primary key, name text);
+        insert into t2 (textid, name) values ('a', 'alpha');
+        insert into t2 (textid, name) values ('b', 'beta');
+        create table if not exists maintable (
+            id integer primary key,
+            fk_integer integer references t1(id),
+            fk_text text references t2(textid)
+        );
+        insert into maintable (id, fk_integer, fk_text) values (1, 1, 'a');
+        insert into maintable (id, fk_integer, fk_text) values (2, 3, 'b'); -- invalid fk_integer
+        insert into maintable (id, fk_integer, fk_text) values (3, 2, 'c'); -- invalid fk_text
+    """
+    )
+    response = await ds.client.get("/db_2214/maintable.csv?_labels=1")
+    assert response.status_code == 200
+    assert response.text == (
+        "id,fk_integer,fk_integer_label,fk_text,fk_text_label\r\n"
+        "1,1,one,a,alpha\r\n"
+        "2,3,,b,beta\r\n"
+        "3,2,two,c,\r\n"
+    )
+
+
 @pytest.mark.asyncio
 async def test_table_csv_blob_columns(ds_client):
     response = await ds_client.get("/fixtures/binary_data.csv")