From a1f383035698da8bf188659390af6e53ffeec940 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 3 Aug 2021 22:20:50 -0700
Subject: [PATCH] --cpu option for datasette publish cloudrun, closes #1420

---
 datasette/publish/cloudrun.py            | 13 +++++-
 docs/datasette-publish-cloudrun-help.txt |  1 +
 tests/test_docs.py                       |  2 +-
 tests/test_publish_cloudrun.py           | 51 ++++++++++++++++--------
 4 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/datasette/publish/cloudrun.py b/datasette/publish/cloudrun.py
index bad223a1..1fabcafd 100644
--- a/datasette/publish/cloudrun.py
+++ b/datasette/publish/cloudrun.py
@@ -36,6 +36,11 @@ def publish_subcommand(publish):
         callback=_validate_memory,
         help="Memory to allocate in Cloud Run, e.g. 1Gi",
     )
+    @click.option(
+        "--cpu",
+        type=click.Choice(["1", "2", "4"]),
+        help="Number of vCPUs to allocate in Cloud Run",
+    )
     @click.option(
         "--apt-get-install",
         "apt_get_extras",
@@ -66,6 +71,7 @@ def publish_subcommand(publish):
         spatialite,
         show_files,
         memory,
+        cpu,
         apt_get_extras,
     ):
         fail_if_publish_binary_not_installed(
@@ -151,8 +157,11 @@ def publish_subcommand(publish):
             image_id = f"gcr.io/{project}/{name}"
             check_call(f"gcloud builds submit --tag {image_id}", shell=True)
         check_call(
-            "gcloud run deploy --allow-unauthenticated --platform=managed --image {} {}{}".format(
-                image_id, service, " --memory {}".format(memory) if memory else ""
+            "gcloud run deploy --allow-unauthenticated --platform=managed --image {} {}{}{}".format(
+                image_id,
+                service,
+                " --memory {}".format(memory) if memory else "",
+                " --cpu {}".format(cpu) if cpu else "",
             ),
             shell=True,
         )
diff --git a/docs/datasette-publish-cloudrun-help.txt b/docs/datasette-publish-cloudrun-help.txt
index 3d05efb6..34481b40 100644
--- a/docs/datasette-publish-cloudrun-help.txt
+++ b/docs/datasette-publish-cloudrun-help.txt
@@ -28,5 +28,6 @@ Options:
   --spatialite                    Enable SpatialLite extension
   --show-files                    Output the generated Dockerfile and metadata.json
   --memory TEXT                   Memory to allocate in Cloud Run, e.g. 1Gi
+  --cpu [1|2|4]                   Number of vCPUs to allocate in Cloud Run
   --apt-get-install TEXT          Additional packages to apt-get install
   --help                          Show this message and exit.
diff --git a/tests/test_docs.py b/tests/test_docs.py
index efd267b9..d0cb036d 100644
--- a/tests/test_docs.py
+++ b/tests/test_docs.py
@@ -50,7 +50,7 @@ def test_help_includes(name, filename):
     # actual has "Usage: cli package [OPTIONS] FILES"
     # because it doesn't know that cli will be aliased to datasette
     expected = expected.replace("Usage: datasette", "Usage: cli")
-    assert expected == actual
+    assert expected == actual, "Run python update-docs-help.py to fix this"
 
 
 @pytest.fixture(scope="session")
diff --git a/tests/test_publish_cloudrun.py b/tests/test_publish_cloudrun.py
index 47f59d72..9c8c38cf 100644
--- a/tests/test_publish_cloudrun.py
+++ b/tests/test_publish_cloudrun.py
@@ -105,17 +105,28 @@ def test_publish_cloudrun(mock_call, mock_output, mock_which, tmp_path_factory):
 @mock.patch("datasette.publish.cloudrun.check_output")
 @mock.patch("datasette.publish.cloudrun.check_call")
 @pytest.mark.parametrize(
-    "memory,should_fail",
+    "memory,cpu,expected_gcloud_args",
     [
-        ["1Gi", False],
-        ["2G", False],
-        ["256Mi", False],
-        ["4", True],
-        ["GB", True],
+        ["1Gi", None, "--memory 1Gi"],
+        ["2G", None, "--memory 2G"],
+        ["256Mi", None, "--memory 256Mi"],
+        ["4", None, None],
+        ["GB", None, None],
+        [None, 1, "--cpu 1"],
+        [None, 2, "--cpu 2"],
+        [None, 3, None],
+        [None, 4, "--cpu 4"],
+        ["2G", 4, "--memory 2G --cpu 4"],
     ],
 )
-def test_publish_cloudrun_memory(
-    mock_call, mock_output, mock_which, memory, should_fail, tmp_path_factory
+def test_publish_cloudrun_memory_cpu(
+    mock_call,
+    mock_output,
+    mock_which,
+    memory,
+    cpu,
+    expected_gcloud_args,
+    tmp_path_factory,
 ):
     mock_output.return_value = "myproject"
     mock_which.return_value = True
@@ -123,22 +134,30 @@ def test_publish_cloudrun_memory(
     os.chdir(tmp_path_factory.mktemp("runner"))
     with open("test.db", "w") as fp:
         fp.write("data")
-    result = runner.invoke(
-        cli.cli,
-        ["publish", "cloudrun", "test.db", "--service", "test", "--memory", memory],
-    )
-    if should_fail:
+    args = ["publish", "cloudrun", "test.db", "--service", "test"]
+    if memory:
+        args.extend(["--memory", memory])
+    if cpu:
+        args.extend(["--cpu", str(cpu)])
+    result = runner.invoke(cli.cli, args)
+    if expected_gcloud_args is None:
         assert 2 == result.exit_code
         return
     assert 0 == result.exit_code
     tag = f"gcr.io/{mock_output.return_value}/datasette"
+    expected_call = (
+        "gcloud run deploy --allow-unauthenticated --platform=managed"
+        " --image {} test".format(tag)
+    )
+    if memory:
+        expected_call += " --memory {}".format(memory)
+    if cpu:
+        expected_call += " --cpu {}".format(cpu)
     mock_call.assert_has_calls(
         [
             mock.call(f"gcloud builds submit --tag {tag}", shell=True),
             mock.call(
-                "gcloud run deploy --allow-unauthenticated --platform=managed --image {} test --memory {}".format(
-                    tag, memory
-                ),
+                expected_call,
                 shell=True,
             ),
         ]