WIP primary keys not null

refs #2807

datasette/views/table_create_alter.py

@@ -772,6 +772,14 @@ def _primary_key_value(columns):
     return tuple(columns)
 
 
+def _primary_key_columns(pk, pks):
+    if pks:
+        return list(pks)
+    if pk:
+        return [pk]
+    return []
+
+
 def _default_expression_sql(default_expr):
     return DEFAULT_EXPR_SQL[default_expr]
 
@@ -871,6 +879,8 @@ class TableCreateView(BaseView):
                 return _error(["pk cannot be changed for existing table"])
             pks = actual_pks
 
+        pk_columns = _primary_key_columns(pk, pks)
+
         initial_schema = None
         if table_exists:
             initial_schema = await db.execute_fn(
@@ -881,11 +891,41 @@ class TableCreateView(BaseView):
             db_for_write = sqlite_utils.Database(conn)
             table = db_for_write[table_name]
             if rows:
+                row_columns = set()
+                pk_columns_set = set(pk_columns)
+                insert_kwargs = {
+                    "pk": pks or pk,
+                    "ignore": ignore,
+                    "replace": replace,
+                    "alter": alter,
+                }
+                if not table_exists and pk_columns:
+                    row_columns = {key for row in rows for key in row}
+                    if pk_columns_set.issubset(row_columns):
+                        insert_kwargs["not_null"] = pk_columns
                 table.insert_all(
-                    rows, pk=pks or pk, ignore=ignore, replace=replace, alter=alter
+                    rows,
+                    **insert_kwargs,
                 )
+                if (
+                    not table_exists
+                    and pk_columns
+                    and not pk_columns_set.issubset(row_columns)
+                ):
+                    table.transform(
+                        not_null={column: True for column in pk_columns},
+                    )
             else:
-                not_null = [column.name for column in columns if column.not_null]
+                column_definitions = {column.name: column.type for column in columns}
+                if pk and pk not in column_definitions:
+                    column_definitions = {pk: "integer", **column_definitions}
+                not_null = [
+                    column.name
+                    for column in columns
+                    if column.not_null or column.name in pk_columns
+                ]
+                if pk and pk not in not_null:
+                    not_null.insert(0, pk)
                 defaults = {}
                 for column in columns:
                     if "default_expr" in column.model_fields_set:
@@ -900,7 +940,7 @@ class TableCreateView(BaseView):
                             db_for_write, column.default
                         )
                 table.create(
-                    {column.name: column.type for column in columns},
+                    column_definitions,
                     pk=pks or pk,
                     foreign_keys=create_request.foreign_keys,
                     not_null=not_null or None,
@@ -1187,6 +1227,7 @@ class TableAlterView(BaseView):
                 defaults = {}
                 column_order = None
                 pk = SQLITE_UTILS_DEFAULT
+                pk_columns = []
                 add_foreign_keys = []
                 drop_foreign_keys = []
                 foreign_keys = None
@@ -1239,6 +1280,7 @@ class TableAlterView(BaseView):
                         drop.add(args.name)
                     elif operation.op == "set_primary_key":
                         pk = _primary_key_value(args.columns)
+                        pk_columns = args.columns
                     elif operation.op == "reorder_columns":
                         column_order = args.columns
                     elif operation.op == "add_foreign_key":
@@ -1281,6 +1323,8 @@ class TableAlterView(BaseView):
                         )
                     )
                     if should_transform:
+                        for column in pk_columns:
+                            not_null[column] = True
                         table.transform(
                             types=types or None,
                             rename=rename or None,

tests/test_api_write.py

@@ -909,6 +909,43 @@ async def test_alter_table_operations(ds_write):
     assert event.after_schema == data["schema"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "table,pk_columns",
+    (
+        ("alter_single_pk", ["id"]),
+        ("alter_compound_pk", ["tenant", "id"]),
+    ),
+)
+async def test_alter_table_primary_keys_are_not_null(ds_write, table, pk_columns):
+    token = write_token(ds_write, permissions=["at"])
+    db = ds_write.get_database("data")
+    await db.execute_write(
+        "create table {} (id integer, tenant text, title text)".format(
+            escape_sqlite(table)
+        )
+    )
+
+    response = await ds_write.client.post(
+        "/data/{}/-/alter".format(table),
+        json={
+            "operations": [
+                {"op": "set_primary_key", "args": {"columns": pk_columns}},
+            ]
+        },
+        headers=_headers(token),
+    )
+
+    assert response.status_code == 200, response.text
+    columns = (
+        await db.execute(
+            "select * from pragma_table_info(?) where pk > 0 order by pk", [table]
+        )
+    ).dicts()
+    assert [column["name"] for column in columns] == pk_columns
+    assert [column["notnull"] for column in columns] == [1] * len(pk_columns)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "default_expr,minimum_value,expected_schema",
@@ -1624,7 +1661,7 @@ async def test_drop_table(ds_write, scenario):
                 "table_api_url": "http://localhost/data/one.json",
                 "schema": (
                     "CREATE TABLE [one] (\n"
-                    "   [id] INTEGER PRIMARY KEY,\n"
+                    "   [id] INTEGER PRIMARY KEY NOT NULL,\n"
                     "   [title] TEXT,\n"
                     "   [score] INTEGER,\n"
                     "   [weight] FLOAT,\n"
@@ -1661,7 +1698,7 @@ async def test_drop_table(ds_write, scenario):
                 "table_api_url": "http://localhost/data/two.json",
                 "schema": (
                     "CREATE TABLE [two] (\n"
-                    "   [id] INTEGER PRIMARY KEY,\n"
+                    "   [id] INTEGER PRIMARY KEY NOT NULL,\n"
                     "   [title] TEXT,\n"
                     "   [score] FLOAT\n"
                     ")"
@@ -1690,7 +1727,7 @@ async def test_drop_table(ds_write, scenario):
                 "table_api_url": "http://localhost/data/three.json",
                 "schema": (
                     "CREATE TABLE [three] (\n"
-                    "   [id] INTEGER PRIMARY KEY,\n"
+                    "   [id] INTEGER PRIMARY KEY NOT NULL,\n"
                     "   [title] TEXT,\n"
                     "   [score] FLOAT\n"
                     ")"
@@ -1734,8 +1771,9 @@ async def test_drop_table(ds_write, scenario):
                 "table_url": "http://localhost/data/five",
                 "table_api_url": "http://localhost/data/five.json",
                 "schema": (
-                    "CREATE TABLE [five] (\n   [type] TEXT,\n   [key] INTEGER,\n"
+                    "CREATE TABLE [five] (\n   [type] TEXT NOT NULL,\n"
-                    "   [title] TEXT,\n   PRIMARY KEY ([type], [key])\n)"
+                    "   [key] INTEGER NOT NULL,\n   [title] TEXT,\n"
+                    "   PRIMARY KEY ([type], [key])\n)"
                 ),
                 "row_count": 1,
             },
@@ -2027,6 +2065,81 @@ async def test_create_table(
     assert [e.name for e in events] == expected_events
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "body,pk_columns",
+    (
+        (
+            {
+                "table": "pk_from_columns",
+                "columns": [
+                    {"name": "id", "type": "integer"},
+                    {"name": "title", "type": "text"},
+                ],
+                "pk": "id",
+            },
+            ["id"],
+        ),
+        (
+            {
+                "table": "compound_pk_from_columns",
+                "columns": [
+                    {"name": "tenant", "type": "text"},
+                    {"name": "id", "type": "integer"},
+                    {"name": "title", "type": "text"},
+                ],
+                "pks": ["tenant", "id"],
+            },
+            ["tenant", "id"],
+        ),
+        (
+            {
+                "table": "pk_omitted_from_columns",
+                "columns": [
+                    {"name": "title", "type": "text"},
+                ],
+                "pk": "id",
+            },
+            ["id"],
+        ),
+        (
+            {
+                "table": "pk_from_rows",
+                "rows": [{"id": 1, "title": "Row 1"}],
+                "pk": "id",
+            },
+            ["id"],
+        ),
+        (
+            {
+                "table": "compound_pk_from_rows",
+                "row": {"tenant": "datasette", "id": 1, "title": "Row 1"},
+                "pks": ["tenant", "id"],
+            },
+            ["tenant", "id"],
+        ),
+    ),
+)
+async def test_create_table_primary_keys_are_not_null(ds_write, body, pk_columns):
+    token = write_token(ds_write)
+    response = await ds_write.client.post(
+        "/data/-/create",
+        json=body,
+        headers=_headers(token),
+    )
+
+    assert response.status_code == 201, response.text
+    db = ds_write.get_database("data")
+    columns = (
+        await db.execute(
+            "select * from pragma_table_info(?) where pk > 0 order by pk",
+            [body["table"]],
+        )
+    ).dicts()
+    assert [column["name"] for column in columns] == pk_columns
+    assert [column["notnull"] for column in columns] == [1] * len(pk_columns)
+
+
 @pytest.mark.asyncio
 async def test_create_table_with_foreign_key(ds_write):
     token = write_token(ds_write)