diff --git a/datasette/publish/cloudrun.py b/datasette/publish/cloudrun.py index ff483bbb..91408287 100644 --- a/datasette/publish/cloudrun.py +++ b/datasette/publish/cloudrun.py @@ -2,6 +2,7 @@ from datasette import hookimpl import click import json import os +import re from subprocess import check_call, check_output from .common import ( @@ -30,6 +31,11 @@ def publish_subcommand(publish): is_flag=True, help="Output the generated Dockerfile and metadata.json", ) + @click.option( + "--memory", + callback=_validate_memory, + help="Memory to allocate in Cloud Run, e.g. 1Gi", + ) def cloudrun( files, metadata, @@ -52,6 +58,7 @@ def publish_subcommand(publish): service, spatialite, show_files, + memory, ): fail_if_publish_binary_not_installed( "gcloud", "Google Cloud", "https://cloud.google.com/sdk/" @@ -127,8 +134,8 @@ def publish_subcommand(publish): image_id = "gcr.io/{project}/{name}".format(project=project, name=name) check_call("gcloud builds submit --tag {}".format(image_id), shell=True) check_call( - "gcloud run deploy --allow-unauthenticated --platform=managed --image {} {}".format( - image_id, service, + "gcloud run deploy --allow-unauthenticated --platform=managed --image {} {}{}".format( + image_id, service, " --memory {}".format(memory) if memory else "" ), shell=True, ) @@ -150,3 +157,9 @@ def get_existing_services(): } for service in services ] + + +def _validate_memory(ctx, param, value): + if re.match(r"^\d+(Gi|G|Mi|M)$", value) is None: + raise click.BadParameter("--memory should be a number then Gi/G/Mi/M e.g 1Gi") + return value diff --git a/tests/test_publish_cloudrun.py b/tests/test_publish_cloudrun.py index b035ca86..272dd68e 100644 --- a/tests/test_publish_cloudrun.py +++ b/tests/test_publish_cloudrun.py @@ -1,6 +1,7 @@ from click.testing import CliRunner from datasette import cli from unittest import mock +import pytest import json @@ -98,6 +99,43 @@ def test_publish_cloudrun(mock_call, mock_output, mock_which): ) +@mock.patch("shutil.which") +@mock.patch("datasette.publish.cloudrun.check_output") +@mock.patch("datasette.publish.cloudrun.check_call") +@pytest.mark.parametrize( + "memory,should_fail", + [["1Gi", False], ["2G", False], ["256Mi", False], ["4", True], ["GB", True],], +) +def test_publish_cloudrun_memory( + mock_call, mock_output, mock_which, memory, should_fail +): + mock_output.return_value = "myproject" + mock_which.return_value = True + runner = CliRunner() + with runner.isolated_filesystem(): + open("test.db", "w").write("data") + result = runner.invoke( + cli.cli, + ["publish", "cloudrun", "test.db", "--service", "test", "--memory", memory], + ) + if should_fail: + assert 2 == result.exit_code + return + assert 0 == result.exit_code + tag = "gcr.io/{}/datasette".format(mock_output.return_value) + mock_call.assert_has_calls( + [ + mock.call("gcloud builds submit --tag {}".format(tag), shell=True), + mock.call( + "gcloud run deploy --allow-unauthenticated --platform=managed --image {} test --memory {}".format( + tag, memory + ), + shell=True, + ), + ] + ) + + @mock.patch("shutil.which") @mock.patch("datasette.publish.cloudrun.check_output") @mock.patch("datasette.publish.cloudrun.check_call")