diff --git a/README.md b/README.md index 4920056..6f280e5 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,7 @@ python generate_custom_image.py \ be escaped, consider encoding the value, then decode it back in the customization script. See more information about VM metadata on https://cloud.google.com/sdk/gcloud/reference/compute/instances/create. +* **--licenses**: Comma-separated list of URIs to license resources. #### Overriding cluster properties with a custom image diff --git a/custom_image_utils/args_parser.py b/custom_image_utils/args_parser.py index b3f48d3..9c8b77d 100644 --- a/custom_image_utils/args_parser.py +++ b/custom_image_utils/args_parser.py @@ -219,6 +219,13 @@ def parse_args(args): instance. This value may need to be increased if your init script generates a lot of output on stdout. If not specified, the default value of 300 seconds will be used.""") + parser.add_argument( + "--licenses", + type=str, + required=False, + default="", + help= + """(Optional) Comma-separated list of URIs to license resources (e.g. 'license-url-a,license-url-b').""") parser.add_argument( "--dry-run", action="store_true", diff --git a/custom_image_utils/shell_script_generator.py b/custom_image_utils/shell_script_generator.py index c7dd8fc..a7632d0 100644 --- a/custom_image_utils/shell_script_generator.py +++ b/custom_image_utils/shell_script_generator.py @@ -119,7 +119,8 @@ --source-disk-zone={zone} \ --source-disk={image_name}-install \ {storage_location_flag} \ - --family={family} + --family={family} \ + {licenses} touch /tmp/{run_id}/image_created }} @@ -172,6 +173,9 @@ def _init_args(self, args): self.args[ "storage_location_flag"] = "--storage-location={storage_location}".format( **self.args) if self.args["storage_location"] else "" + self.args[ + "licenses"] = "--licenses='{licenses}'".format( + **self.args) if self.args["licenses"] else "" metadata_flag_template = ( "--metadata=shutdown-timer-in-sec={shutdown_timer_in_sec}," "custom-sources-path={custom_sources_path}") diff --git a/tests/test_args_parser.py b/tests/test_args_parser.py index c726ba6..a11d0b7 100644 --- a/tests/test_args_parser.py +++ b/tests/test_args_parser.py @@ -49,6 +49,7 @@ def test_minimal_required_args(self): family="'dataproc-custom-image'", gcs_bucket="'{}'".format(gcs_bucket), image_name="'{}'".format(image_name), + licenses="''", machine_type="'n1-standard-1'", network="'{}'".format(''), no_external_ip="False", @@ -74,6 +75,7 @@ def test_optional_args(self): family = 'debian9' gcs_bucket = 'gs://my-bucket' image_name = 'my-image' + licenses = 'https://www.googleapis.com/compute/v1/projects/license/my-project/global/licenses/my-license' machine_type = 'n1-standard-4' network = 'my-network' no_external_ip = True @@ -96,6 +98,7 @@ def test_optional_args(self): '--family', family, '--gcs-bucket', gcs_bucket, '--image-name', image_name, + '--licenses', licenses, '--machine-type', machine_type, '--network', network, '--no-external-ip', @@ -122,6 +125,7 @@ def test_optional_args(self): family="'{}'".format(family), gcs_bucket="'{}'".format(gcs_bucket), image_name="'{}'".format(image_name), + licenses="'{}'".format(licenses), machine_type="'{}'".format(machine_type), metadata="'{}'".format(metadata), network="'{}'".format(network), @@ -150,6 +154,7 @@ def _make_expected_result( family, gcs_bucket, image_name, + licenses, machine_type, metadata, network, @@ -175,6 +180,7 @@ def _make_expected_result( "family={}, " "gcs_bucket={}, " "image_name={}, " + "licenses={}, " "machine_type={}, " "metadata={}, " "network={}, " @@ -199,6 +205,7 @@ def _make_expected_result( family, gcs_bucket, image_name, + licenses, machine_type, metadata, network, diff --git a/tests/test_shell_script_generator.py b/tests/test_shell_script_generator.py index f59b7bc..9749062 100644 --- a/tests/test_shell_script_generator.py +++ b/tests/test_shell_script_generator.py @@ -88,7 +88,7 @@ fi echo 'Creating custom image.' - gcloud compute images create my-image --project=my-project --source-disk-zone=us-west1-a --source-disk=my-image-install --storage-location=us-east1 --family=debian9 + gcloud compute images create my-image --project=my-project --source-disk-zone=us-west1-a --source-disk=my-image-install --storage-location=us-east1 --family=debian9 --licenses='https://www.googleapis.com/compute/v1/projects/my-project/global/licenses/my-license' touch /tmp/custom-image-my-image-20190611-160823/image_created } @@ -122,7 +122,8 @@ def test_generate_shell_script(self): 'project_id': 'my-project', 'storage_location': 'us-east1', 'shutdown_timer_in_sec': 500, - 'base_image_family': 'projects/my-dataproc-project/global/images/family/debian-10' + 'base_image_family': 'projects/my-dataproc-project/global/images/family/debian-10', + 'licenses': 'https://www.googleapis.com/compute/v1/projects/my-project/global/licenses/my-license' } script = shell_script_generator.Generator().generate(args)