Prototype of add/edit foreign keys feature, refs #7

datasette_edit_schema/__init__.py

@@ -3,6 +3,10 @@
 from datasette.utils import sqlite3
 from urllib.parse import quote_plus, unquote_plus
 import sqlite_utils
+from .utils import get_primary_keys, potential_foreign_keys
+
+# Don't attempt to detect foreign keys on tables larger than this:
+FOREIGN_KEY_DETECTION_LIMIT = 10_000
 
 
 @hookimpl
@@ -213,6 +217,11 @@ def transform_the_table(conn):
 
             return Response.redirect(request.path)
 
+        if formdata.get("action") == "update_foreign_keys":
+            return await update_foreign_keys(
+                request, datasette, database, table, formdata
+            )
+
         if "delete_table" in formdata:
             return await delete_table(request, datasette, database, table)
         elif "add_column" in formdata:
@@ -222,14 +231,21 @@ def transform_the_table(conn):
         else:
             return Response.html("Unknown operation", status=400)
 
-    def get_columns_and_schema(conn):
+    def get_columns_and_schema_and_fks(conn):
         t = sqlite_utils.Database(conn)[table]
+        pks = set(t.pks)
         columns = [
-            {"name": column, "type": dtype} for column, dtype in t.columns_dict.items()
+            {"name": column, "type": dtype, "is_pk": column in pks}
+            for column, dtype in t.columns_dict.items()
         ]
-        return columns, t.schema
+        return columns, t.schema, t.foreign_keys
 
-    columns, schema = await database.execute_fn(get_columns_and_schema)
+    columns, schema, foreign_keys = await database.execute_fn(
+        get_columns_and_schema_and_fks
+    )
+    foreign_keys_by_column = {}
+    for fk in foreign_keys:
+        foreign_keys_by_column.setdefault(fk.column, []).append(fk)
 
     columns_display = [
         {
@@ -239,6 +255,105 @@ def get_columns_and_schema(conn):
         for c in columns
     ]
 
+    # To detect potential foreign keys we need (table, column) for the
+    # primary keys on every other table
+    other_primary_keys = [
+        pair for pair in await database.execute_fn(get_primary_keys) if pair[0] != table
+    ]
+    integer_primary_keys = [
+        (pair[0], pair[1]) for pair in other_primary_keys if pair[2] is int
+    ]
+    string_primary_keys = [
+        (pair[0], pair[1]) for pair in other_primary_keys if pair[2] is str
+    ]
+
+    column_foreign_keys = [
+        {
+            "name": column["name"],
+            "foreign_key": foreign_keys_by_column.get(column["name"])[0]
+            if foreign_keys_by_column.get(column["name"])
+            else None,
+            "suggestions": [],
+            "options": integer_primary_keys
+            if column["type"] is int
+            else string_primary_keys,
+        }
+        for column in columns
+        if not column["is_pk"]
+    ]
+    # Only scan for potential foreign keys if there are less than 10,000
+    # rows - since execute_fn() does not yet support time limits
+    if (
+        await database.execute(
+            'select count(*) from (select 1 from "{}" limit {})'.format(
+                table, FOREIGN_KEY_DETECTION_LIMIT
+            )
+        )
+    ).single_value() < FOREIGN_KEY_DETECTION_LIMIT:
+        potential_fks = await database.execute_fn(
+            lambda conn: potential_foreign_keys(
+                conn,
+                table,
+                [c["name"] for c in columns if not c["is_pk"]],
+                other_primary_keys,
+            )
+        )
+        for info in column_foreign_keys:
+            info["suggestions"] = potential_fks.get(info["name"], [])
+
+    # Add 'options' to those
+    for info in column_foreign_keys:
+        options = []
+        seen = set()
+        info["html_options"] = options
+        # Reshuffle so suggestions are at the top
+        if info["foreign_key"]:
+            options.append(
+                {
+                    "name": "{}.{} (current)".format(
+                        info["foreign_key"].other_table,
+                        info["foreign_key"].other_column,
+                    ),
+                    "value": "{}.{}".format(
+                        info["foreign_key"].other_table,
+                        info["foreign_key"].other_column,
+                    ),
+                    "selected": True,
+                }
+            )
+            seen.add(
+                "{}:{}".format(
+                    info["foreign_key"].other_table, info["foreign_key"].other_column
+                )
+            )
+        # Now add suggestions
+        for suggested_table, suggested_column in info["suggestions"]:
+            if not (
+                info["foreign_key"]
+                and info["foreign_key"].other_column == suggested_column
+            ):
+                options.append(
+                    {
+                        "name": "{}.{} (suggested)".format(
+                            suggested_table, suggested_column
+                        ),
+                        "value": "{}.{}".format(suggested_table, suggested_column),
+                        "selected": False,
+                    }
+                )
+                seen.add("{}:{}".format(suggested_table, suggested_column))
+                info["suggested"] = "{}.{}".format(suggested_table, suggested_column)
+        # And the rest
+        for rest_table, rest_column in info["options"]:
+            if "{}:{}".format(rest_table, rest_column) not in seen:
+                options.append(
+                    {
+                        "name": "{}.{}".format(rest_table, rest_column),
+                        "value": "{}.{}".format(rest_table, rest_column),
+                        "selected": False,
+                    }
+                )
+
     return Response.html(
         await datasette.render_template(
             "edit_schema_table.html",
@@ -251,6 +366,8 @@ def get_columns_and_schema(conn):
                     {"name": TYPE_NAMES[value], "value": value}
                     for value in TYPES.values()
                 ],
+                "foreign_keys": foreign_keys,
+                "column_foreign_keys": column_foreign_keys,
             },
             request=request,
         )
@@ -350,3 +467,52 @@ async def rename_table(request, datasette, database, table, formdata):
     return Response.redirect(
         "/-/edit-schema/{}/{}".format(quote_plus(database.name), quote_plus(new_name))
     )
+
+
+async def update_foreign_keys(request, datasette, database, table, formdata):
+    new_fks = {
+        key[3:]: value
+        for key, value in formdata.items()
+        if key.startswith("fk.") and value.strip()
+    }
+    existing_fks = {
+        fk.column: fk.other_table + "." + fk.other_column
+        for fk in await database.execute_fn(
+            lambda conn: sqlite_utils.Database(conn)[table].foreign_keys
+        )
+    }
+    if new_fks == existing_fks:
+        datasette.add_message(request, "No changes to foreign keys", datasette.WARNING)
+        return Response.redirect(request.path)
+
+    # Need that in (column, other_table, other_column) format
+    fks = []
+    for column, other_table_and_column in new_fks.items():
+        split = other_table_and_column.split(".")
+        fks.append(
+            (
+                column,
+                split[0],
+                split[1],
+            )
+        )
+
+    # Update foreign keys
+    def run(conn):
+        db = sqlite_utils.Database(conn)
+        with conn:
+            db[table].transform(foreign_keys=fks)
+
+    await database.execute_write_fn(run, block=True)
+    summary = ", ".join("{} → {}.{}".format(*fk) for fk in fks)
+    if summary:
+        message = "Foreign keys updated{}".format(
+            " to {}".format(summary) if summary else ""
+        )
+    else:
+        message = "Foreign keys removed"
+    datasette.add_message(
+        request,
+        message,
+    )
+    return Response.redirect(request.path)

datasette_edit_schema/templates/edit_schema_table.html

@@ -120,6 +120,34 @@ <h2>Add a column</h2>
     <input type="submit" value="Add column">
 </form>
 
+<h2>Update foreign key relationships</h2>
+
+<p>Configure foreign keys on columns so Datasette can link related tables together.</p>
+
+<style type="text/css">
+table.foreign-key-options td {
+    white-space: normal;
+}
+</style>
+
+<form action="{{ base_url }}-/edit-schema/{{ database.name|quote_plus }}/{{ table|quote_plus }}" method="post">
+    <input type="hidden" name="csrftoken" value="{{ csrftoken() }}">
+    <input type="hidden" name="action" value="update_foreign_keys">
+    <table class="foreign-key-options">
+    {% for column in column_foreign_keys %}
+      <tr>
+        <td><label for="fk.{{ column.name }}">{{ column.name }}</label></td>
+        <td><select id="fk.{{ column.name }}" name="fk.{{ column.name }}"><option value="">-- {% if not column.suggested and not column.foreign_key %}no suggestions{% else %}none{% endif %} --</option>
+            {% for option in column.html_options %}<option value="{{ option.value }}" {% if option.selected %} selected="selected"{% endif %}>{{ option.name }}</option>{% endfor %}
+            </select>
+            {% if column.suggested %}<p style="margin: 0; font-size: 0.8em">Suggested: {{ column.suggested }}</p>{% endif %}
+            </td>
+      </tr>
+    {% endfor %}
+    </table>
+    <input type="submit" value="Update foreign keys">
+</form>
+
 <h2>Delete table</h2>
 
 <form action="{{ base_url }}-/edit-schema/{{ database.name|quote_plus }}/{{ table|quote_plus }}" method="post">

datasette_edit_schema/utils.py

@@ -0,0 +1,47 @@
+import sqlite_utils
+
+
+def get_primary_keys(conn):
+    db = sqlite_utils.Database(conn)
+    primary_keys = []
+    for table in db.tables:
+        if "_fts_" in table.name:
+            continue
+        pks = table.pks
+        if pks == ["rowid"]:
+            continue
+        if len(pks) != 1:
+            continue
+        pk = pks[0]
+        # Is that a str or int?
+        pk_type = table.columns_dict[pk]
+        if pk_type in (str, int):
+            primary_keys.append((table.name, pk, pk_type))
+    return primary_keys
+
+
+def potential_foreign_keys(conn, table_name, columns, other_table_pks):
+    potentials = {}
+    cursor = conn.cursor()
+    for column in columns:
+        potentials[column] = []
+        for other_table, other_column, _ in other_table_pks:
+            query = """
+                select "{table}"."{column}"
+                from "{table}"
+                where not exists (
+                    select 1
+                    from "{other_table}"
+                    where "{table}"."{column}" = "{other_table}"."{other_column}"
+                )
+                limit 1;
+            """.format(
+                table=table_name,
+                column=column,
+                other_table=other_table,
+                other_column=other_column,
+            )
+            cursor.execute(query)
+            if cursor.fetchone() is None:
+                potentials[column].append((other_table, other_column))
+    return potentials

setup.py

@@ -25,7 +25,7 @@ def get_long_description():
     entry_points={"datasette": ["edit_schema = datasette_edit_schema"]},
     install_requires=[
         "datasette>=0.63",
-        "sqlite-utils>=3.10",
+        "sqlite-utils>=3.35",
     ],
     extras_require={"test": ["pytest", "pytest-asyncio"]},
     package_data={"datasette_edit_schema": ["templates/*.html", "static/*.js"]},