Added support for licenses

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
 

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",

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}")

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,

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)