Azure: add image filter

This lets the user specify an image via attribute filters, similar
to AWS.

Change-Id: Ie86314c0a90a66550d288704c443fce41e356ad9

doc/source/azure.rst

@@ -311,18 +311,65 @@ section of the configuration.
            long-standing issue with ``ansible_shell_type`` in combination
            with ``become``
 
+      .. attr:: image-filter
+         :type: dict
+
+         Specifies a private image to use via filters.  Either this field,
+         :attr:`providers.[azure].cloud-images.image-reference`, or
+         :attr:`providers.[azure].cloud-images.image-id` must be
+         provided.
+
+         If a filter is provided, Nodepool will list all of the images
+         in the provider's resource group and reduce the list using
+         the supplied filter.  All items specified in the filter must
+         match in order for an image to match.  If more than one image
+         matches, the images are sorted by name and the last one
+         matches.
+
+         Example:
+
+         .. code-block:: yaml
+
+            cloud-images:
+              - name: image-by-name
+                image-filter:
+                  name: test-image
+              - name: image-by-tag
+                image-filter:
+                  tags:
+                    foo: bar
+
+         The following filters are available:
+
+         .. attr:: name
+            :type: str
+
+            The name of the image.
+
+         .. attr:: location
+            :type: str
+
+            The location of the image.
+
+         .. attr:: tags
+            :type: dict
+
+            The image tags.
+
       .. attr:: image-id
          :type: str
 
-         Specifies a private image to use.  Either this field or
+         Specifies a private image to use by ID.  Either this field,
-         :attr:`providers.[azure].cloud-images.image-reference` must be
+         :attr:`providers.[azure].cloud-images.image-reference`, or
+         :attr:`providers.[azure].cloud-images.image-filter` must be
          provided.
 
       .. attr:: image-reference
          :type: dict
 
-         Specifies a public image to use.  Either this field or
+         Specifies a public image to use.  Either this field,
-         :attr:`providers.[azure].cloud-images.image-id` must be
+         :attr:`providers.[azure].cloud-images.image-id`, or
+         :attr:`providers.[azure].cloud-images.image-filter` must be
          provided.
 
          .. attr:: sku

nodepool/driver/azure/adapter.py

@@ -174,6 +174,7 @@ class AzureCreateStateMachine(statemachine.StateMachine):
         self.retries = retries
         self.attempts = 0
         self.image_external_id = image_external_id
+        self.image_reference = None
         self.metadata = metadata
         self.tags = label.tags.copy() or {}
         self.tags.update(metadata)
@@ -203,6 +204,12 @@ class AzureCreateStateMachine(statemachine.StateMachine):
     def advance(self):
         if self.state == self.START:
             self.external_id = self.hostname
+
+            # Find an appropriate image if filters were provided
+            if self.label.cloud_image and self.label.cloud_image.image_filter:
+                self.image_reference = self.adapter._getImageFromFilter(
+                    self.label.cloud_image.image_filter)
+
             if self.label.pool.public_ipv4:
                 self.public_ipv4 = self.adapter._createPublicIPAddress(
                     self.tags, self.hostname, self.ip_sku, 'IPv4',
@@ -234,8 +241,9 @@ class AzureCreateStateMachine(statemachine.StateMachine):
             self.nic = self.adapter._refresh(self.nic)
             if self.adapter._succeeded(self.nic):
                 self.vm = self.adapter._createVirtualMachine(
-                    self.label, self.image_external_id, self.tags,
+                    self.label, self.image_external_id,
-                    self.hostname, self.nic)
+                    self.image_reference, self.tags, self.hostname,
+                    self.nic)
                 self.state = self.VM_CREATING
             else:
                 return
@@ -647,13 +655,20 @@ class AzureAdapter(statemachine.Adapter):
         with self.rate_limiter:
             return self.azul.virtual_machines.list(self.resource_group)
 
-    def _createVirtualMachine(self, label, image_external_id, tags,
+    def _createVirtualMachine(self, label, image_external_id,
-                              hostname, nic):
+                              image_reference, tags, hostname, nic):
         if image_external_id:
+            # This is a diskimage
             image = label.diskimage
             remote_image = self._getImage(image_external_id)
             image_reference = {'id': remote_image['id']}
+        elif image_reference:
+            # This is a cloud image with aser supplied image-filter;
+            # we already found the reference.
+            image = label.cloud_image
         else:
+            # This is a cloud image with a user-supplied reference or
+            # id.
             image = label.cloud_image
             if label.cloud_image.image_reference:
                 image_reference = label.cloud_image.image_reference
@@ -741,3 +756,24 @@ class AzureAdapter(statemachine.Adapter):
     def _listImages(self):
         with self.rate_limiter:
             return self.azul.images.list(self.resource_group)
+
+    def _getImageFromFilter(self, image_filter):
+        images = self._listImages()
+        images = [i for i in images
+                  if i['properties']['provisioningState'] == 'Succeeded']
+        if 'name' in image_filter:
+            images = [i for i in images
+                      if i['name'] == image_filter['name']]
+        if 'location' in image_filter:
+            images = [i for i in images
+                      if i['location'] == image_filter['location']]
+        if 'tags' in image_filter:
+            for k, v in image_filter['tags'].items():
+                images = [i for i in images if i['tags'].get(k) == v]
+        images = sorted(images, key=lambda i: i['name'])
+        if not images:
+            raise Exception("Unable to find image matching filter: %s",
+                            image_filter)
+        image = images[-1]
+        self.log.debug("Found image matching filter: %s", image)
+        return {'id': image['id']}

nodepool/driver/azure/config.py

@@ -37,6 +37,7 @@ class AzureProviderCloudImage(ConfigValue):
         # TODO(corvus): remove zuul_public_key
         self.key = image.get('key', zuul_public_key)
         self.image_reference = image.get('image-reference')
+        self.image_filter = image.get('image-filter')
         self.image_id = image.get('image-id')
         self.python_path = image.get('python-path')
         self.shell_type = image.get('shell-type')
@@ -48,7 +49,8 @@ class AzureProviderCloudImage(ConfigValue):
     @property
     def external_name(self):
         '''Human readable version of external.'''
-        return self.image_id or self.image_reference or self.name
+        return (self.image_id or self.image_reference or
+                self.image_filter or self.name)
 
     @staticmethod
     def getSchema():
@@ -59,6 +61,12 @@ class AzureProviderCloudImage(ConfigValue):
             v.Required('offer'): str,
         }
 
+        azure_image_filter = {
+            'location': str,
+            'name': str,
+            'tags': dict,
+        }
+
         return v.All({
             v.Required('name'): str,
             v.Required('username'): str,
@@ -68,14 +76,16 @@ class AzureProviderCloudImage(ConfigValue):
             'key': str,
             v.Exclusive('image-reference', 'spec'): azure_image_reference,
             v.Exclusive('image-id', 'spec'): str,
+            v.Exclusive('image-filter', 'spec'): azure_image_filter,
             'connection-type': str,
             'connection-port': int,
             'python-path': str,
             'shell-type': str,
         }, {
             v.Required(
-                v.Any('image-reference', 'image-id'),
+                v.Any('image-reference', 'image-id', 'image-filter'),
-                msg='Provide either "image-reference" or "image-id" keys'
+                msg=('Provide either "image-reference", '
+                     '"image-filter", or "image-id" keys')
             ): object,
             object: object,
         })

nodepool/tests/fixtures/azure.yaml

@@ -19,6 +19,10 @@ labels:
     min-ready: 0
   - name: windows-generate
     min-ready: 0
+  - name: image-by-name
+    min-ready: 0
+  - name: image-by-tag
+    min-ready: 0
 
 providers:
   - name: azure
@@ -54,6 +58,15 @@ providers:
           offer: WindowsServer
         username: foobar
         generate-password: True
+      - name: image-by-name
+        username: zuul
+        image-filter:
+          name: test1
+      - name: image-by-tag
+        username: zuul
+        image-filter:
+          tags:
+            foo: bar
     pools:
       - name: main
         max-servers: 10
@@ -71,6 +84,14 @@ providers:
               systemPurpose: CI
             user-data: "This is the user data"
             custom-data: "This is the custom data"
+          - name: image-by-name
+            cloud-image: image-by-name
+            hardware-profile:
+              vm-size: Standard_B1ls
+          - name: image-by-tag
+            cloud-image: image-by-tag
+            hardware-profile:
+              vm-size: Standard_B1ls
           - name: windows-password
             cloud-image: windows-password
             hardware-profile:

nodepool/tests/unit/test_driver_azure.py

@@ -23,6 +23,34 @@ from nodepool.driver.statemachine import StateMachineProvider
 from . import fake_azure
 
 
+def make_image(name, tags):
+    return {
+        'name': name,
+        'id': ('/subscriptions/c35cf7df-ed75-4c85-be00-535409a85120/'
+               'resourceGroups/nodepool/providers/Microsoft.Compute/'
+               f'images/{name}'),
+        'type': 'Microsoft.Compute/images',
+        'location': 'eastus',
+        'tags': tags,
+        'properties': {
+            'storageProfile': {
+                'osDisk': {
+                    'osType': 'Linux',
+                    'osState': 'Generalized',
+                    'diskSizeGB': 1,
+                    'blobUri': 'https://example.net/nodepoolstorage/img.vhd',
+                    'caching': 'ReadWrite',
+                    'storageAccountType': 'Standard_LRS'
+                },
+                'dataDisks': [],
+                'zoneResilient': False
+            },
+            'provisioningState': 'Succeeded',
+            'hyperVGeneration': 'V1'
+        }
+    }
+
+
 class TestDriverAzure(tests.DBTestCase):
     log = logging.getLogger("nodepool.TestDriverAzure")
 
@@ -140,6 +168,76 @@ class TestDriverAzure(tests.DBTestCase):
             "/resourceGroups/nodepool/providers/Microsoft.Compute"
             "/images/test-image-1234")
 
+    def test_azure_image_filter_name(self):
+        self.fake_azure.crud['Microsoft.Compute/images'].items.append(
+            make_image('test1', {'foo': 'bar'}))
+        self.fake_azure.crud['Microsoft.Compute/images'].items.append(
+            make_image('test2', {}))
+        self.fake_azure.crud['Microsoft.Compute/images'].items.append(
+            make_image('test3', {'foo': 'bar'}))
+
+        configfile = self.setup_config(
+            'azure.yaml',
+            auth_path=self.fake_azure.auth_file.name)
+        pool = self.useNodepool(configfile, watermark_sleep=1)
+        pool.start()
+        req = zk.NodeRequest()
+        req.state = zk.REQUESTED
+        req.node_types.append('image-by-name')
+
+        self.zk.storeNodeRequest(req)
+        req = self.waitForNodeRequest(req)
+
+        self.assertEqual(req.state, zk.FULFILLED)
+        self.assertNotEqual(req.nodes, [])
+        node = self.zk.getNode(req.nodes[0])
+        self.assertEqual(node.allocated_to, req.id)
+        self.assertEqual(node.state, zk.READY)
+        self.assertIsNotNone(node.launcher)
+        self.assertEqual(node.connection_type, 'ssh')
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            requests[0]['properties']['storageProfile']
+            ['imageReference']['id'],
+            "/subscriptions/c35cf7df-ed75-4c85-be00-535409a85120"
+            "/resourceGroups/nodepool/providers/Microsoft.Compute"
+            "/images/test1")
+
+    def test_azure_image_filter_tag(self):
+        self.fake_azure.crud['Microsoft.Compute/images'].items.append(
+            make_image('test1', {'foo': 'bar'}))
+        self.fake_azure.crud['Microsoft.Compute/images'].items.append(
+            make_image('test2', {}))
+        self.fake_azure.crud['Microsoft.Compute/images'].items.append(
+            make_image('test3', {'foo': 'bar'}))
+
+        configfile = self.setup_config(
+            'azure.yaml',
+            auth_path=self.fake_azure.auth_file.name)
+        pool = self.useNodepool(configfile, watermark_sleep=1)
+        pool.start()
+        req = zk.NodeRequest()
+        req.state = zk.REQUESTED
+        req.node_types.append('image-by-tag')
+
+        self.zk.storeNodeRequest(req)
+        req = self.waitForNodeRequest(req)
+
+        self.assertEqual(req.state, zk.FULFILLED)
+        self.assertNotEqual(req.nodes, [])
+        node = self.zk.getNode(req.nodes[0])
+        self.assertEqual(node.allocated_to, req.id)
+        self.assertEqual(node.state, zk.READY)
+        self.assertIsNotNone(node.launcher)
+        self.assertEqual(node.connection_type, 'ssh')
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            requests[0]['properties']['storageProfile']
+            ['imageReference']['id'],
+            "/subscriptions/c35cf7df-ed75-4c85-be00-535409a85120"
+            "/resourceGroups/nodepool/providers/Microsoft.Compute"
+            "/images/test3")
+
     def test_azure_windows_image_password(self):
         configfile = self.setup_config(
             'azure.yaml',

releasenotes/notes/azure-image-filter-d44182130ffb7225.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Added support for filtering Azure images.  Use the
+    :attr:`providers.[azure].cloud-images.image-filter` setting to
+    specify a private image using filters (tags, for example).