WIP /db/table/-/upsert, refs #1878

datasette/app.py

@@ -40,7 +40,7 @@
     PermissionsDebugView,
     MessagesDebugView,
 )
-from .views.table import TableView, TableInsertView, TableDropView
+from .views.table import TableView, TableInsertView, TableUpsertView, TableDropView
 from .views.row import RowView, RowDeleteView, RowUpdateView
 from .renderer import json_renderer
 from .url_builder import Urls
@@ -1292,6 +1292,10 @@ def add_route(view, regex):
             TableInsertView.as_view(self),
             r"/(?P<database>[^\/\.]+)/(?P<table>[^\/\.]+)/-/insert$",
         )
+        add_route(
+            TableUpsertView.as_view(self),
+            r"/(?P<database>[^\/\.]+)/(?P<table>[^\/\.]+)/-/upsert$",
+        )
         add_route(
             TableDropView.as_view(self),
             r"/(?P<database>[^\/\.]+)/(?P<table>[^\/\.]+)/-/drop$",

datasette/views/special.py

@@ -316,21 +316,37 @@ async def example_links(self, request):
                     request.actor, "insert-row", (name, table)
                 ):
                     pks = await db.primary_keys(table)
-                    table_links.append(
-                        {
-                            "path": self.ds.urls.table(name, table) + "/-/insert",
-                            "method": "POST",
-                            "label": "Insert rows into {}".format(table),
-                            "json": {
-                                "rows": [
-                                    {
-                                        column: None
-                                        for column in await db.table_columns(table)
-                                        if column not in pks
-                                    }
-                                ]
+                    table_links.extend(
+                        [
+                            {
+                                "path": self.ds.urls.table(name, table) + "/-/insert",
+                                "method": "POST",
+                                "label": "Insert rows into {}".format(table),
+                                "json": {
+                                    "rows": [
+                                        {
+                                            column: None
+                                            for column in await db.table_columns(table)
+                                            if column not in pks
+                                        }
+                                    ]
+                                },
                             },
-                        }
+                            {
+                                "path": self.ds.urls.table(name, table) + "/-/upsert",
+                                "method": "POST",
+                                "label": "Upsert rows into {}".format(table),
+                                "json": {
+                                    "rows": [
+                                        {
+                                            column: None
+                                            for column in await db.table_columns(table)
+                                            if column not in pks
+                                        }
+                                    ]
+                                },
+                            },
+                        ]
                     )
                 if await self.ds.permission_allowed(
                     request.actor, "drop-table", (name, table)

datasette/views/table.py

@@ -1074,9 +1074,15 @@ class TableInsertView(BaseView):
     def __init__(self, datasette):
         self.ds = datasette
 
-    async def _validate_data(self, request, db, table_name):
+    async def _validate_data(self, request, db, table_name, pks, upsert):
         errors = []
 
+        pks_list = []
+        if isinstance(pks, str):
+            pks_list = [pks]
+        else:
+            pks_list = list(pks)
+
         def _errors(errors):
             return None, errors, {}
 
@@ -1135,6 +1141,15 @@ def _errors(errors):
         # Validate columns of each row
         columns = set(await db.table_columns(table_name))
         for i, row in enumerate(rows):
+            if upsert:
+                # It MUST have the primary key
+                missing_pks = [pk for pk in pks_list if pk not in row]
+                if missing_pks:
+                    errors.append(
+                        'Row {} is missing primary key column(s): "{}"'.format(
+                            i, '", "'.join(missing_pks)
+                        )
+                    )
             invalid_columns = set(row.keys()) - columns
             if invalid_columns:
                 errors.append(
@@ -1146,7 +1161,7 @@ def _errors(errors):
             return _errors(errors)
         return rows, errors, extras
 
-    async def post(self, request):
+    async def post(self, request, upsert=False):
         try:
             resolved = await self.ds.resolve_table(request)
         except NotFound as e:
@@ -1164,34 +1179,44 @@ async def post(self, request):
             request.actor, "insert-row", resource=(database_name, table_name)
         ):
             return _error(["Permission denied"], 403)
-        rows, errors, extras = await self._validate_data(request, db, table_name)
+
+        pks = await db.primary_keys(table_name)
+
+        rows, errors, extras = await self._validate_data(
+            request, db, table_name, pks, upsert
+        )
         if errors:
             return _error(errors, 400)
 
         ignore = extras.get("ignore")
         replace = extras.get("replace")
 
         should_return = bool(extras.get("return", False))
-        # Insert rows
-        def insert_rows(conn):
+
+        def insert_or_upsert_rows(conn):
             table = sqlite_utils.Database(conn)[table_name]
+            kwargs = {}
+            if upsert:
+                kwargs["pk"] = pks[0] if len(pks) == 1 else pks
+            else:
+                kwargs = {"ignore": ignore, "replace": replace}
             if should_return:
                 rowids = []
+                method = table.upsert if upsert else table.insert
                 for row in rows:
-                    rowids.append(
-                        table.insert(row, ignore=ignore, replace=replace).last_rowid
-                    )
+                    rowids.append(method(row, **kwargs).last_rowid)
                 return list(
                     table.rows_where(
                         "rowid in ({})".format(",".join("?" for _ in rowids)),
                         rowids,
                     )
                 )
             else:
-                table.insert_all(rows, ignore=ignore, replace=replace)
+                method_all = table.upsert_all if upsert else table.insert_all
+                method_all(rows, **kwargs)
 
         try:
-            rows = await db.execute_write_fn(insert_rows)
+            rows = await db.execute_write_fn(insert_or_upsert_rows)
         except Exception as e:
             return _error([str(e)])
         result = {"ok": True}
@@ -1200,6 +1225,13 @@ def insert_rows(conn):
         return Response.json(result, status=201)
 
 
+class TableUpsertView(TableInsertView):
+    name = "table-upsert"
+
+    async def post(self, request):
+        return await super().post(request, upsert=True)
+
+
 class TableDropView(BaseView):
     name = "table-drop"