From 2b20957b1869b5a2213960d2e3a67188f42d2a2f Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 25 Jan 2022 19:19:39 -0800
Subject: [PATCH] Better validation for upsert --pk, closes #390

---
 docs/cli-reference.rst |   3 +-
 sqlite_utils/cli.py    | 118 +++++++++++++++++++++--------------------
 tests/test_cli.py      |  17 ++++++
 3 files changed, 79 insertions(+), 59 deletions(-)

diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst
index 583f7e2..eee907e 100644
--- a/docs/cli-reference.rst
+++ b/docs/cli-reference.rst
@@ -253,8 +253,6 @@ See :ref:`cli_upsert`.
       incoming record has a primary key that matches an existing record the existing
       record will be updated.
 
-      The --pk option is required.
-
       Example:
 
           echo '[
@@ -264,6 +262,7 @@ See :ref:`cli_upsert`.
 
     Options:
       --pk TEXT                 Columns to use as the primary key, e.g. id
+                                [required]
       --flatten                 Flatten nested JSON objects, so {"a": {"b": 1}}
                                 becomes {"a_b": 1}
       --nl                      Expect newline-delimited JSON
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 350f9e3..65a108f 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -813,59 +813,65 @@ def import_options(fn):
     return fn
 
 
-def insert_upsert_options(fn):
-    for decorator in reversed(
-        (
-            click.argument(
-                "path",
-                type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
-                required=True,
-            ),
-            click.argument("table"),
-            click.argument("file", type=click.File("rb"), required=True),
-            click.option(
-                "--pk", help="Columns to use as the primary key, e.g. id", multiple=True
-            ),
-        )
-        + _import_options
-        + (
-            click.option(
-                "--batch-size", type=int, default=100, help="Commit every X records"
-            ),
-            click.option(
-                "--alter",
-                is_flag=True,
-                help="Alter existing table to add any missing columns",
-            ),
-            click.option(
-                "--not-null",
-                multiple=True,
-                help="Columns that should be created as NOT NULL",
-            ),
-            click.option(
-                "--default",
-                multiple=True,
-                type=(str, str),
-                help="Default value that should be set for a column",
-            ),
-            click.option(
-                "-d",
-                "--detect-types",
-                is_flag=True,
-                envvar="SQLITE_UTILS_DETECT_TYPES",
-                help="Detect types for columns in CSV/TSV data",
-            ),
-            click.option(
-                "--analyze",
-                is_flag=True,
-                help="Run ANALYZE at the end of this operation",
-            ),
-            load_extension_option,
-            click.option("--silent", is_flag=True, help="Do not show progress bar"),
-        )
-    ):
-        fn = decorator(fn)
-    return fn
+def insert_upsert_options(*, require_pk=False):
+    def inner(fn):
+        for decorator in reversed(
+            (
+                click.argument(
+                    "path",
+                    type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
+                    required=True,
+                ),
+                click.argument("table"),
+                click.argument("file", type=click.File("rb"), required=True),
+                click.option(
+                    "--pk",
+                    help="Columns to use as the primary key, e.g. id",
+                    multiple=True,
+                    required=require_pk,
+                ),
+            )
+            + _import_options
+            + (
+                click.option(
+                    "--batch-size", type=int, default=100, help="Commit every X records"
+                ),
+                click.option(
+                    "--alter",
+                    is_flag=True,
+                    help="Alter existing table to add any missing columns",
+                ),
+                click.option(
+                    "--not-null",
+                    multiple=True,
+                    help="Columns that should be created as NOT NULL",
+                ),
+                click.option(
+                    "--default",
+                    multiple=True,
+                    type=(str, str),
+                    help="Default value that should be set for a column",
+                ),
+                click.option(
+                    "-d",
+                    "--detect-types",
+                    is_flag=True,
+                    envvar="SQLITE_UTILS_DETECT_TYPES",
+                    help="Detect types for columns in CSV/TSV data",
+                ),
+                click.option(
+                    "--analyze",
+                    is_flag=True,
+                    help="Run ANALYZE at the end of this operation",
+                ),
+                load_extension_option,
+                click.option("--silent", is_flag=True, help="Do not show progress bar"),
+            )
+        ):
+            fn = decorator(fn)
+        return fn
+
+    return inner
 
 
 def insert_upsert_implementation(
@@ -1060,7 +1066,7 @@ def _find_variables(tb, vars):
 
 
 @cli.command()
-@insert_upsert_options
+@insert_upsert_options()
 @click.option(
     "--ignore", is_flag=True, default=False, help="Ignore records if pk already exists"
 )
@@ -1168,7 +1174,7 @@ def insert(
 
 
 @cli.command()
-@insert_upsert_options
+@insert_upsert_options(require_pk=True)
 def upsert(
     path,
     table,
@@ -1201,8 +1207,6 @@ def upsert(
     an incoming record has a primary key that matches an existing record
     the existing record will be updated.
 
-    The --pk option is required.
-
     Example:
 
     \b
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 60c95b0..21b7945 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -922,6 +922,23 @@ def test_upsert(db_path, tmpdir):
     ]
 
 
+def test_upsert_pk_required(db_path, tmpdir):
+    json_path = str(tmpdir / "dogs.json")
+    db = Database(db_path)
+    insert_dogs = [
+        {"id": 1, "name": "Cleo", "age": 4},
+        {"id": 2, "name": "Nixie", "age": 4},
+    ]
+    open(json_path, "w").write(json.dumps(insert_dogs))
+    result = CliRunner().invoke(
+        cli.cli,
+        ["upsert", db_path, "dogs", json_path],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 2
+    assert "Error: Missing option '--pk'" in result.output
+
+
 def test_upsert_analyze(db_path, tmpdir):
     db = Database(db_path)
     db["rows"].insert({"id": 1, "foo": "x", "n": 3}, pk="id")