Extending Azure CLI

Andrei Zhozhin

| 10 minutes

Why to extend Azure CLI?

Azure CLI already has a lot of extensions to cover interactions with various services in Azure cloud. It also available on all platforms (Linux, Windows, MacOS) so same tool could be used in various scenarios. Many developers already have experience working with Azure CLI so if we add new “extension” it would feel natural and would not require learning new tool from scratch. But it would require to “design” groups and commands so they would feel natural and logical.

Azure CLI contains code to perform authentication so extension could reuse existing tokens or request new using core library functions.

In this article I would look at the process of creation of new extension to work with PowerBI service. At the time of this article creation there are no official PowerBI extensions available from Microsoft, so it might be useful.

Every Azure CLI command look like the following:

az [ group ] [ subgroup ] [ command ] {parameters}

So we would introduce new group, multiple subgroups to cover various PowerBI objects and multiple commands to implement something like CRUD (Create, Read, Update, Delete) operations for these objects.

Azure CLI modules vs extensions

There are two repositories which contains Azure CLI itself + modules and Extensions for it.

Some modules live in Azure CLI repository and released together with Azure CLI package. Other extensions are stored in its own repository to allow rapid development and independent release cycle with indirect dependency on Azure CLI.

Azure CLI bundled extensions could be found here: azure-cli/azure/cli/command_modules

Extensions (Azure CLI extensions repo)

Pros Cons
Released on its own Explicit install required az extension add
Rapid development possible Can be broken by azure-cli-core changes
Experimental UX is allowed
Leverage CLI code generation

Modules (part of Azure CLI repo)

Pros Cons
Comes automatically with Azure CLI Strictly tied to Azure CLI schedule
Unlikely to be broken by azure-cli-core changes Strictly controlled, experimental functionality might be not allowed

Our choice is extension as we are planning to implement experimental functionality and we want even to host it separately from azure-cli-extensions

Workflow

Service -> Python SDK -> Azure CLI

  1. We have some cloud service that expose API. It is PowerBI service in our case. And we have swagger file for it.
  2. We want to have client for this service to interact with it from our extension. We would use PowerBI swagger to generate Python client.
  3. And finally we want to have command line commands that would use service client to interact with the service.

Generating PowerBI service client

We want to interact with PowerBI service via REST API, so we either use existing python library for that or generate client based on Swagger definition.

Here is a Swagger file from Microsoft’s PowerBI C# client: https://github.com/microsoft/PowerBI-CSharp/blob/master/sdk/swaggers/swagger.json

PowerBI swagger file is massive and pretty complex, it also has some “preview” endpoints that sometimes cause problems during code generation. So we need to make some adjustments to this file to make it work properly in our use-case:

  • Add missing "produces": ["application/json"] mime-types to certain operations (commit)
  • Fix format for string based enum for columnId (commit)
  • Fix "comsumes": ["application/json"] for PostImport and PostImportInGroup (commit)
  • Fix escaping for username description to avoid build errors in python code generator
  • Set parameter location to method for generated code to have not-polluted configuration object (commit)

Complete swagger file with all amendments could be found here: azhozhin/powerbi-swagger

Create the following powerbi.yaml file

input-file: swaggers\swagger.json
use-extension: 
    "@autorest/python" : "6.9.7"
python:
    output-folder: src
    add-credentials: true

and run generation

autorest powerbi.yaml

Generated client would be available in src folder. Now we need to move generated client into src\pbi\vendored_sdks.

Azure CLI extension structure

Every Azure CLI extensions should be placed into corresponding folder in /src Our extension would be named pbi to avoid intersection with powerbi which is already occupied by powerbidedicated extension powerbidedicated/azext_powerbidedicated/commands.py

So our extension pbi would have the following structure (folder src/pbi):

.
├── __init__.py                 # empty file
├── azext_pbi
│   ├── __init__.py             # entry point
│   ├── _client_factory.py      # powerbi client factory
│   ├── _help.py                # help/docs
│   ├── _params.py              # parameter registration
│   ├── azext_metadata.json     # additional 
│   ├── commands.py             # command registration
│   ├── custom.py               # command implementations
│   ├── tests                   
│   │   ├── __init__.py         # empty file
│   │   ├── recordings          # recorded web responses
│   │   ├── test_*_scenario.py  # test scenarios
│   ├── vendored_sdks
│   │   ├── __init__.py         # empty file
│   │   └── power_bi_client     # generated client files (from swagger)
│   └── version.py              # version file
└── setup.py                    # python package file

I’ll cover creation of Azure CLI extension from scratch, but now there is new developer tool that could help generating boilerplate and speedup extension development - azdev

Adding groups and commands

Lets add three commands distributed into 2 groups.

group commands
pbi version
pbi workspace list, show

in commands.py it would look like that:

def load_command_table(self, args):  
    with self.command_group('pbi') as g:  
        g.custom_command('version', 'pbi_version')  
  
    with self.command_group('pbi workspace') as g:  
        g.custom_command('list', 'pbi_workspace_list')  
        g.custom_command('show', 'pbi_workspace_show')

based on entry point registration (azext_pbi/__init__.py) all these methods should be in custom.py

class PowerBICLICommandsLoader(AzCommandsLoader):  
    def __init__(self, cli_ctx=None):  
        pbi_custom = CliCommandType(
	        operations_tmpl='azext_pbi.custom#{}')  # all custom groups/commands
        super(PowerBICLICommandsLoader, self).__init__(cli_ctx=cli_ctx, 
	        custom_command_type=pbi_custom)
...

And adding commands implementation in custom.py:

from azure.core.exceptions import ResourceNotFoundError  
from ._client_factory import _powerbi_client_factory  
from .version import VERSION  
  

# this method return only version as object
def pbi_version():  
    return {  
        'version': VERSION  
    }  
  
# return list of all groups/workspaces visible to user
def pbi_workspace_list(cmd):  
    client = _powerbi_client_factory(cmd.cli_ctx)  
    groups = client.groups.get_groups()  
    return groups['value']  
  
# return details only for selected workspace
def pbi_workspace_show(cmd, workspace):  
    client = _powerbi_client_factory(cmd.cli_ctx)  
    groups = client.groups.get_groups()  
    for d in groups['value']:  
        if d['name'] == workspace:  
            return d  
    raise ResourceNotFoundError(f'Workspace "{workspace}" is not found')

At this stage we should be able to compile our extension

python setup.py bdist_wheel

And install extension from local folder

az extension add --source ./dist/azure_pbi-0.0.1-py3-none-any

Now we can try to run command that just return version to validate that everything works fine.

az pbi version

{
  "version": "0.0.1"
}

If you authenticate to azure

az login --allow-no-subscription

Then if your account has access to powerbi workspaces you can get list of workspaces (I’ve changed workspace GUIDs)

az pbi workspace list
[
  {
    "id": "00000000-0000-0000-0000-000000000001",
    "isOnDedicatedCapacity": false,
    "isReadOnly": false,
    "name": "FirstWorkspace",
    "type": "Workspace"
  },
  {
    "id": "00000000-0000-0000-0000-000000000002",
    "isOnDedicatedCapacity": false,
    "isReadOnly": false,
    "name": "SecondWorkspace",
    "type": "Workspace"
  },
  {
    "id": "00000000-0000-0000-0000-000000000003",
    "isOnDedicatedCapacity": false,
    "isReadOnly": false,
    "name": "Admin monitoring",
    "type": "AdminInsights"
  }
]

and querying individual workspace by name

az pbi workspace show --workspace "Admin monitoring"
{
  "id": "00000000-0000-0000-0000-000000000003",
  "isOnDedicatedCapacity": false,
  "isReadOnly": false,
  "name": "Admin monitoring",
  "type": "AdminInsights"
}

Adding documentation/help messages

If we add --help option to our new command we would not see any description, only list of arguments

az pbi workspace show --help

Command
    az pbi workspace show

Arguments
    --workspace [Required]

We want to add some description to parameter as well as command description.

Adding parameter docs first. In _params.py we want to register new type for workspace parameter:

from knack.arguments import CLIArgumentType  
  
workspace_type = CLIArgumentType(
    help='Workspace name',
    options_list=['--workspace', '-w']
)

def load_arguments(self, _):
    # specify context where we are attaching argument information
    with self.argument_context('pbi workspace show') as c:
        c.argument('workspace', workspace_type)

And adding command description in _help.py. Help message should have special YAML formatting.

helps['pbi workspace show'] = """  
    type: command    
    short-summary: Show workspace details by name    
    examples:      
       - name: Show workspace details        
         text: |-            
               az pbi workspace show --workspace "My Workspace Name"
"""

if we run previous command again:

az pbi workspace show --help

Command
    az pbi workspace show : Show workspace details by name.

Arguments
    --workspace -w [Required] : Workspace name.

...

Examples
    Show workspace details
        az pbi workspace show --workspace "My Workspace Name"

Now we have documentation for groups, commands, and parameters

Testing

Azure CLI already have extensive test infrastructure so we can run commands that requires network interactions and these interactions would be recorded for consequent runs.

Testing pbi version command (which does not require network interaction) so we use standard unittest approach:

import unittest  
  
from azext_pbi.custom import pbi_version  
from azext_pbi.version import VERSION  
  
  
class VersionTest(unittest.TestCase):  
  
    def test_version(self):  
        actual = pbi_version()  
        assert actual['version'] == VERSION

Test method would execute pbi_version method and compare it with expected VERSION constant.

Now lets do more complicated tests that require PowerBI service interaction. Our test class should be derived from ScenarioTest class which implements testing harness.


from azure.cli.testsdk import ScenarioTest  
  
  
class WorkspaceScenarioTest(ScenarioTest):  
  
    def test_workspace(self):  
        workspace_name = 'Admin monitoring'  
        # first call to PowerBI API 
        self.cmd('az pbi workspace list').checks = [  
            self.check('length([])', 3)  
        ]  
  
        self.kwargs.update({  
            "name": workspace_name,  
        })  

        # second call to PowerBI API
        self.cmd('az pbi workspace show --workspace "{name}"', checks=[  
            self.check('isOnDedicatedCapacity', False),  
            self.check('isReadOnly', False),  
            self.check('name', workspace_name),  
            self.check('type', 'AdminInsights'),  
        ])

All checks (self.check) by default using JMESPath expression to validate output (which should be JSON). So the following statement

self.check('length([])')

Would return length of resulting array from command call. We do not care about internals of every returned item yet, we only check that list is really returning list of items with expected length.

Now we want to check particular named workspace and its internals.

Resulting JSON of the second command az pbi workspace show --workspace "{name}":

{
  "id": "00000000-0000-0000-0000-000000000003",
  "isOnDedicatedCapacity": false,
  "isReadOnly": false,
  "name": "Admin monitoring",
  "type": "AdminInsights"
}

And four checks for the second command call:

self.check('isOnDedicatedCapacity', False),
self.check('isReadOnly', False),
self.check('name', workspace_name),
self.check('type', 'AdminInsights'),

Would validate that output structure has expected values in their corresponding properties.

In our test scenario we have two steps:

  • get list of all workspace (az pbi workspace list)
  • getting workspace by name (az pbi workspace show --workspace "{name}")

On every call to PowerBI REST API azure test sdk would record a network call (src/pbi/azext_pbi/tests/recordings/test_workspace.yaml) and reply it next time, so no external call would be performed. It is very useful and allow testing without network interaction which should be very fast.

Release cycle

After we are done with implementation and testing it is time to prepare package to be released so other people can use it. If you want to publish your extension to central Azure CLI extensions repository please follow standard release requirements.

There are two components we need to care about:

  1. Wheel (*.whl) package with out extension - this one need to be published to PYPI repository
  2. And extension index index.json - this one need to be published somewhere accessible to everyone (even on github pages)

Publishing Wheel

I would use TestPYPI to not pollute central python index with my half-baked package.

Building Wheel package

python -m build --wheel

Validating it with twine

twine check dist/*

Publishing to testpypi

twine upload --repository testpypi dist/*

Publishing index.json

I would generate and publish index.json to github pages. Index file is special index that we need to prepare. Original file with centrally managed extensions could be found here

I’ve prepared python script that could generate index.json based on dist folder. It is not production ready but would demonstrate how to create one. Let’s run it in src/pbi

python ../../scripts/generate_index.py

Now we have index.json with generated reference to PYPI package that we’ve published on previous step.

Lets put this file into pages branch and configure github to host our static site from root (/)

After that file should be available : https://azhozhin.github.io/azure-cli-my-extensions/index.json

And we finally can install extension adding index explicitly.

az extension add --name azure-pbi \
                 --index https://azhozhin.github.io/azure-cli-my-extensions/index.json

Conclusion

We have implemented all steps required to create new Azure CLI extension:

  1. Generated PowerBI python client based on swagger
  2. Created two command groups (pbi and pbi workspace)
  3. Created three commands
    1. pbi version
    2. pbi workspace list
    3. pbi workspace show
  4. Added documentation for groups and commands
  5. Added documentation for argument --workspace -w <name>
  6. Added unit test for pbi version and integration tests for pbi workspace list and pbi workspace show commands
  7. Prepared wheel(*.whl) package and index.json
  8. Published package to TestPYPI registry
  9. Published index.json to static site so others could install our extension

Our extension is pretty simple at this stage but now we can easily extend it to cover more functionality that is provided by PowerBI API.

Full source code available here: azhozhin/azure-cli-my-extensions

References

Previous post

BigData London 2022

Related content