From 23be5be1dc94299b987f0f47eedda15fb45854f0 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 16 Mar 2024 17:05:27 -0700
Subject: [PATCH] create-table now supports multiple --pk, refs #620

---
 docs/cli.rst        | 2 ++
 sqlite_utils/cli.py | 6 +++---
 tests/test_cli.py   | 8 ++++++++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/docs/cli.rst b/docs/cli.rst
index ed60544..0254ab1 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -1909,6 +1909,8 @@ This will create a table called ``mytable`` with two columns - an integer ``id``
 
 You can pass as many column-name column-type pairs as you like. Valid types are ``integer``, ``text``, ``float`` and ``blob``.
 
+Pass ``--pk`` more than once for a compound primary key that covers multiple columns.
+
 You can specify columns that should be NOT NULL using ``--not-null colname``. You can specify default values for columns using ``--default colname defaultvalue``.
 
 .. code-block:: bash
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 33f661d..53f8ee0 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -1489,7 +1489,7 @@ def create_database(path, enable_wal, init_spatialite, load_extension):
 )
 @click.argument("table")
 @click.argument("columns", nargs=-1, required=True)
-@click.option("--pk", help="Column to use as primary key")
+@click.option("pks", "--pk", help="Column to use as primary key", multiple=True)
 @click.option(
     "--not-null",
     multiple=True,
@@ -1532,7 +1532,7 @@ def create_table(
     path,
     table,
     columns,
-    pk,
+    pks,
     not_null,
     default,
     fk,
@@ -1581,7 +1581,7 @@ def create_table(
             )
     db[table].create(
         coltypes,
-        pk=pk,
+        pk=pks[0] if len(pks) == 1 else pks,
         not_null=not_null,
         defaults=dict(default),
         foreign_keys=fk,
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 6a9c581..bc04fd4 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1180,6 +1180,14 @@ def test_upsert_alter(db_path, tmpdir):
             ["age", "integer", "--default", "age", "3"],
             ("CREATE TABLE [t] (\n" "   [age] INTEGER DEFAULT '3'\n" ")"),
         ),
+        # Compound primary key
+        (
+            ["category", "text", "name", "text", "--pk", "category", "--pk", "name"],
+            (
+                "CREATE TABLE [t] (\n   [category] TEXT,\n   [name] TEXT,\n"
+                "   PRIMARY KEY ([category], [name])\n)"
+            ),
+        ),
     ],
 )
 def test_create_table(args, schema):