From 45e30e9f2b2439043ad74840aea2b05f2d2cbc4a Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 6 Sep 2023 15:29:35 -0700 Subject: [PATCH] Prototype of add/edit foreign keys feature, refs #7 --- datasette_edit_schema/__init__.py | 174 +++++++++++++++++- .../templates/edit_schema_table.html | 28 +++ datasette_edit_schema/utils.py | 47 +++++ setup.py | 2 +- 4 files changed, 246 insertions(+), 5 deletions(-) create mode 100644 datasette_edit_schema/utils.py diff --git a/datasette_edit_schema/__init__.py b/datasette_edit_schema/__init__.py index d895a42..dda1627 100644 --- a/datasette_edit_schema/__init__.py +++ b/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) diff --git a/datasette_edit_schema/templates/edit_schema_table.html b/datasette_edit_schema/templates/edit_schema_table.html index e08356c..e5f09af 100644 --- a/datasette_edit_schema/templates/edit_schema_table.html +++ b/datasette_edit_schema/templates/edit_schema_table.html @@ -120,6 +120,34 @@

Add a column

+

Update foreign key relationships

+ +

Configure foreign keys on columns so Datasette can link related tables together.

+ + + +
+ + + + {% for column in column_foreign_keys %} + + + + + {% endfor %} +
+ {% if column.suggested %}

Suggested: {{ column.suggested }}

{% endif %} +
+ +
+

Delete table

diff --git a/datasette_edit_schema/utils.py b/datasette_edit_schema/utils.py new file mode 100644 index 0000000..648c43d --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index e462537..1e4be19 100644 --- a/setup.py +++ b/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"]},