From b190b87ec6897872ad3111ec694219cb9de98579 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 1 Feb 2025 17:02:49 -0800
Subject: [PATCH] Detect single unique text column in label_column_for_table,
 closes #2458

Also added new tests for label_column_for_table()
---
 datasette/database.py                | 30 ++++++++-
 tests/test_label_column_for_table.py | 97 ++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_label_column_for_table.py

diff --git a/datasette/database.py b/datasette/database.py
index a2e899bc..4a0babfb 100644
--- a/datasette/database.py
+++ b/datasette/database.py
@@ -3,6 +3,7 @@ from collections import namedtuple
 from pathlib import Path
 import janus
 import queue
+import sqlite_utils
 import sys
 import threading
 import uuid
@@ -442,7 +443,33 @@ class Database:
         )
         if explicit_label_column:
             return explicit_label_column
-        column_names = await self.execute_fn(lambda conn: table_columns(conn, table))
+
+        def column_details(conn):
+            # Returns {column_name: (type, is_unique)}
+            db = sqlite_utils.Database(conn)
+            columns = db[table].columns_dict
+            indexes = db[table].indexes
+            details = {}
+            for name in columns:
+                is_unique = any(
+                    index
+                    for index in indexes
+                    if index.columns == [name] and index.unique
+                )
+                details[name] = (columns[name], is_unique)
+            return details
+
+        column_details = await self.execute_fn(column_details)
+        # Is there just one unique column that's text?
+        unique_text_columns = [
+            name
+            for name, (type_, is_unique) in column_details.items()
+            if is_unique and type_ is str
+        ]
+        if len(unique_text_columns) == 1:
+            return unique_text_columns[0]
+
+        column_names = list(column_details.keys())
         # Is there a name or title column?
         name_or_title = [c for c in column_names if c.lower() in ("name", "title")]
         if name_or_title:
@@ -452,6 +479,7 @@ class Database:
             column_names
             and len(column_names) == 2
             and ("id" in column_names or "pk" in column_names)
+            and not set(column_names) == {"id", "pk"}
         ):
             return [c for c in column_names if c not in ("id", "pk")][0]
         # Couldn't find a label:
diff --git a/tests/test_label_column_for_table.py b/tests/test_label_column_for_table.py
new file mode 100644
index 00000000..7667b595
--- /dev/null
+++ b/tests/test_label_column_for_table.py
@@ -0,0 +1,97 @@
+import pytest
+from datasette.database import Database
+from datasette.app import Datasette
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "create_sql,table_name,config,expected_label_column",
+    [
+        # Explicit label_column
+        (
+            "create table t1 (id integer primary key, name text, title text);",
+            "t1",
+            {"t1": {"label_column": "title"}},
+            "title",
+        ),
+        # Single unique text column
+        (
+            "create table t2 (id integer primary key, name2 text unique, title text);",
+            "t2",
+            {},
+            "name2",
+        ),
+        (
+            "create table t3 (id integer primary key, title2 text unique, name text);",
+            "t3",
+            {},
+            "title2",
+        ),
+        # Two unique text columns means it cannot decide on one
+        (
+            "create table t3x (id integer primary key, name2 text unique, title2 text unique);",
+            "t3x",
+            {},
+            None,
+        ),
+        # Name or title column
+        (
+            "create table t4 (id integer primary key, name text);",
+            "t4",
+            {},
+            "name",
+        ),
+        (
+            "create table t5 (id integer primary key, title text);",
+            "t5",
+            {},
+            "title",
+        ),
+        # But not if there are multiple non-unique text that are not called title
+        (
+            "create table t5x (id integer primary key, other1 text, other2 text);",
+            "t5x",
+            {},
+            None,
+        ),
+        (
+            "create table t6 (id integer primary key, Name text);",
+            "t6",
+            {},
+            "Name",
+        ),
+        (
+            "create table t7 (id integer primary key, Title text);",
+            "t7",
+            {},
+            "Title",
+        ),
+        # Two columns, one of which is id
+        (
+            "create table t8 (id integer primary key, content text);",
+            "t8",
+            {},
+            "content",
+        ),
+        (
+            "create table t9 (pk integer primary key, content text);",
+            "t9",
+            {},
+            "content",
+        ),
+    ],
+)
+async def test_label_column_for_table(
+    create_sql, table_name, config, expected_label_column
+):
+    """Test cases for label_column_for_table method"""
+    ds = Datasette()
+    db = ds.add_database(Database(ds, memory_name="test_label_column_for_table"))
+    await db.execute_write_script(create_sql)
+    if config:
+        ds.config = {"databases": {"test_label_column_for_table": {"tables": config}}}
+    actual_label_column = await db.label_column_for_table(table_name)
+    if expected_label_column is None:
+        assert actual_label_column is None
+    else:
+        assert actual_label_column == expected_label_column